@aztec/foundation 0.0.0-test.1 → 0.0.1-commit.b655e406

package/src/json-rpc/server/safe_json_rpc_server.ts

@@ -23,6 +23,21 @@ export type DiagnosticsData = {
 
 export type DiagnosticsMiddleware = (ctx: DiagnosticsData, next: () => Promise<void>) => Promise<void>;
 
+export type SafeJsonRpcServerConfig = {
+  /** Maximum batch size for batched rpc requests */
+  maxBatchSize: number;
+  /** Return an HTTP 200 status code on errors, but include an error object as per the JSON RPC spec */
+  http200OnError: boolean;
+  /** The maximum body size the server will accept */
+  maxBodySizeBytes: string;
+};
+
+const defaultServerConfig: SafeJsonRpcServerConfig = {
+  http200OnError: false,
+  maxBatchSize: 100,
+  maxBodySizeBytes: '50mb',
+};
+
 export class SafeJsonRpcServer {
   /**
    * The HTTP server accepting remote requests.
@@ -30,21 +45,26 @@ export class SafeJsonRpcServer {
    */
   private httpServer?: http.Server;
 
+  private config: SafeJsonRpcServerConfig;
+
   constructor(
-    /** The proxy object to delegate requests to. */
+    /** The proxy object to delegate requests to */
     private readonly proxy: Proxy,
-    /**
-     *  Return an HTTP 200 status code on errors, but include an error object
-     *  as per the JSON RPC spec
-     */
-    private http200OnError = false,
+    config: Partial<SafeJsonRpcServerConfig> = {},
     /** Health check function */
     private readonly healthCheck: StatusCheckFn = () => true,
     /** Additional middlewares */
     private extraMiddlewares: Application.Middleware[] = [],
     /** Logger */
     private log = createLogger('json-rpc:server'),
-  ) {}
+  ) {
+    this.config = { ...defaultServerConfig, ...config };
+
+    // handle empty string
+    if (!this.config.maxBodySizeBytes) {
+      this.config.maxBodySizeBytes = defaultServerConfig.maxBodySizeBytes;
+    }
+  }
 
   public isHealthy(): boolean | Promise<boolean> {
     return this.healthCheck();
@@ -63,20 +83,13 @@ export class SafeJsonRpcServer {
         await next();
       } catch (err: any) {
         const method = (ctx.request.body as any)?.method ?? 'unknown';
-        this.log.warn(`Error in JSON RPC server call ${method}: ${inspect(err)}`);
-        if (err instanceof SyntaxError) {
+        this.log.warn(`Uncaught error in JSON RPC server call ${method}: ${inspect(err)}`);
+        if (err && 'name' in err && err.name === 'BadRequestError') {
           ctx.status = 400;
-          ctx.body = { jsonrpc: '2.0', id: null, error: { code: -32700, message: `Parse error: ${err.message}` } };
-        } else if (err instanceof ZodError) {
-          const message = err.issues.map(e => `${e.message} (${e.path.join('.')})`).join('. ') || 'Validation error';
+          ctx.body = { jsonrpc: '2.0', id: null, error: { code: -32000, message: `Bad request: ${err.message}` } };
+        } else if (err && err instanceof SyntaxError) {
           ctx.status = 400;
-          ctx.body = { jsonrpc: '2.0', id: null, error: { code: -32701, message } };
-        } else if (this.http200OnError) {
-          ctx.body = {
-            jsonrpc: '2.0',
-            id: null,
-            error: { code: err.code || -32600, data: err.data, message: err.message },
-          };
+          ctx.body = { jsonrpc: '2.0', id: null, error: { code: -32700, message: `Parse error: ${err.message}` } };
         } else {
           ctx.status = 500;
           ctx.body = { jsonrpc: '2.0', id: null, error: { code: -32600, message: err.message ?? 'Internal error' } };
@@ -107,7 +120,12 @@ export class SafeJsonRpcServer {
       app.use(middleware);
     }
     app.use(exceptionHandler);
-    app.use(bodyParser({ jsonLimit: '50mb', enableTypes: ['json'], detectJSON: () => true }));
+    app.use(
+      bodyParser({
+        jsonLimit: this.config.maxBodySizeBytes,
+        enableTypes: ['json'],
+      }),
+    );
     app.use(cors());
     app.use(router.routes());
     app.use(router.allowedMethods());
@@ -124,23 +142,93 @@ export class SafeJsonRpcServer {
     const router = new Router({ prefix });
     // "JSON RPC mode" where a single endpoint is used and the method is given in the request body
     router.post('/', async (ctx: Koa.Context) => {
-      const { params = [], jsonrpc, id, method } = ctx.request.body as any;
-      // Fail if not a registered function in the proxy
-      if (typeof method !== 'string' || method === 'constructor' || !this.proxy.hasMethod(method)) {
-        ctx.status = 400;
-        const code = -32601;
-        const message = `Method not found: ${method}`;
-        ctx.body = { jsonrpc, id, error: { code, message } };
+      if (Array.isArray(ctx.request.body)) {
+        if (ctx.request.body.length > this.config.maxBatchSize) {
+          ctx.status = this.config.http200OnError ? 200 : 400;
+          ctx.body = {
+            jsonrpc: '2.0',
+            error: {
+              code: -32600,
+              message: `Batch request exceeds maximum allowed size of ${this.config.maxBatchSize}`,
+            },
+            id: null,
+          };
+          return;
+        }
+        const resp = await this.processBatch(ctx.request.body);
+        if (Array.isArray(resp)) {
+          ctx.status = 200;
+          ctx.body = resp;
+        } else {
+          ctx.status = this.config.http200OnError ? 200 : 400;
+          ctx.body = resp;
+        }
       } else {
-        ctx.status = 200;
-        const result = await this.proxy.call(method, params);
-        ctx.body = { jsonrpc, id, result };
+        const resp = await this.processRequest(ctx.request.body);
+        if ('error' in resp) {
+          ctx.status = this.config.http200OnError ? 200 : 400;
+        }
+
+        ctx.body = resp;
       }
     });
 
     return router;
   }
 
+  private async processBatch(requests: any[]) {
+    if (requests.length === 0) {
+      return { jsonrpc: '2.0', error: { code: -32600, message: 'Invalid Request' }, id: null };
+    }
+    const results = await Promise.allSettled(requests.map(req => this.processRequest(req)));
+    return results.map(res => {
+      if (res.status === 'fulfilled') {
+        return res.value;
+      }
+
+      this.log.warn(`Uncaught error executing request in batch: ${res.reason}.`);
+      return { jsonrpc: '2.0', error: { code: -32600, message: 'Invalid Request' }, id: null };
+    });
+  }
+
+  private async processRequest(request: any) {
+    if (!request || typeof request !== 'object') {
+      return { jsonrpc: '2.0', error: { code: -32600, message: 'Invalid Request' }, id: null };
+    }
+
+    const { params = [], jsonrpc = '2.0', id, method } = request;
+    if (typeof method !== 'string' || !method) {
+      return { jsonrpc: '2.0', id, error: { code: -32600, message: `Invalid request` } };
+    }
+
+    // Fail if not a registered function in the proxy
+    if (typeof method !== 'string' || method === 'constructor' || !this.proxy.hasMethod(method)) {
+      return { jsonrpc, id, error: { code: -32601, message: `Method not found: ${method}` } };
+    } else {
+      try {
+        const result = await this.proxy.call(method, params);
+        return { jsonrpc, id, result };
+      } catch (err: any) {
+        if (err && err instanceof ZodError) {
+          const message = err.issues.map(e => `${e.message} (${e.path.join('.')})`).join('. ') || 'Validation error';
+          return { jsonrpc: '2.0', id, error: { code: -32701, message } };
+        } else if (err) {
+          return {
+            jsonrpc,
+            id,
+            error: { code: -32702, data: err.data, message: err.message },
+          };
+        } else {
+          return {
+            jsonrpc,
+            id,
+            error: { code: -32702, message: 'Error executing request' },
+          };
+        }
+      }
+    }
+  }
+
   /**
    * Start this server with koa.
    * @param port - Port number.
@@ -201,7 +289,10 @@ export class SafeJsonProxy<T extends object = any> implements Proxy {
   private log = createLogger('json-rpc:proxy');
   private schema: ApiSchema;
 
-  constructor(private handler: T, schema: ApiSchemaFor<T>) {
+  constructor(
+    private handler: T,
+    schema: ApiSchemaFor<T>,
+  ) {
     this.schema = schema;
   }
 
@@ -265,7 +356,10 @@ function makeAggregateHealthcheck(namedHandlers: NamespacedApiHandlers, log?: Lo
   return async () => {
     try {
       const results = await Promise.all(
-        Object.entries(namedHandlers).map(([name, [, , healthCheck]]) => [name, healthCheck ? healthCheck() : true]),
+        Object.entries(namedHandlers).map(async ([name, [, , healthCheck]]) => [
+          name,
+          healthCheck ? await healthCheck() : true,
+        ]),
       );
       const failed = results.filter(([_, result]) => !result);
       if (failed.length > 0) {
@@ -280,12 +374,13 @@ function makeAggregateHealthcheck(namedHandlers: NamespacedApiHandlers, log?: Lo
   };
 }
 
-export type SafeJsonRpcServerOptions = {
-  http200OnError: boolean;
-  healthCheck?: StatusCheckFn;
-  log?: Logger;
-  middlewares?: Application.Middleware[];
-};
+export type SafeJsonRpcServerOptions = Partial<
+  SafeJsonRpcServerConfig & {
+    healthCheck: StatusCheckFn;
+    log: Logger;
+    middlewares: Application.Middleware[];
+  }
+>;
 
 /**
  * Creates a single SafeJsonRpcServer from multiple handlers.
@@ -294,24 +389,22 @@ export type SafeJsonRpcServerOptions = {
  */
 export function createNamespacedSafeJsonRpcServer(
   handlers: NamespacedApiHandlers,
-  options: Partial<Omit<SafeJsonRpcServerOptions, 'healthcheck'>> = {
-    log: createLogger('json-rpc:server'),
-  },
+  options: Omit<SafeJsonRpcServerOptions, 'healthcheck'> = {},
 ): SafeJsonRpcServer {
-  const { middlewares, http200OnError, log } = options;
+  const { middlewares, log } = options;
   const proxy = new NamespacedSafeJsonProxy(handlers);
   const healthCheck = makeAggregateHealthcheck(handlers, log);
-  return new SafeJsonRpcServer(proxy, http200OnError, healthCheck, middlewares, log);
+  return new SafeJsonRpcServer(proxy, options, healthCheck, middlewares, log);
 }
 
 export function createSafeJsonRpcServer<T extends object = any>(
   handler: T,
   schema: ApiSchemaFor<T>,
-  options: Partial<SafeJsonRpcServerOptions> = {},
+  options: SafeJsonRpcServerOptions = {},
 ) {
-  const { http200OnError, log, healthCheck, middlewares: extraMiddlewares } = options;
+  const { log, healthCheck, middlewares: extraMiddlewares } = options;
   const proxy = new SafeJsonProxy(handler, schema);
-  return new SafeJsonRpcServer(proxy, http200OnError, healthCheck, extraMiddlewares, log);
+  return new SafeJsonRpcServer(proxy, options, healthCheck, extraMiddlewares, log);
 }
 
 /**
@@ -326,7 +419,7 @@ export function createStatusRouter(getCurrentStatus: StatusCheckFn, apiPrefix =
     let ok: boolean;
     try {
       ok = (await getCurrentStatus()) === true;
-    } catch (err) {
+    } catch {
       ok = false;
     }
 
@@ -364,7 +457,7 @@ export async function startHttpRpcServer(
 
   const { promise, resolve } = promiseWithResolvers<void>();
   const listenPort = options.port ? (typeof options.port === 'string' ? parseInt(options.port) : options.port) : 0;
-  httpServer.listen(listenPort, options.host, () => resolve());
+  httpServer.listen({ host: options.host, port: listenPort, reuseAddress: true }, () => resolve());
 
   // Wait until listen callback is called
   if (!options.noWait) {

package/src/json-rpc/test/integration.ts

@@ -20,7 +20,7 @@ export type JsonRpcTestContext<T extends object> = {
 export async function createJsonRpcTestSetup<T extends object>(
   handler: T,
   schema: ApiSchemaFor<T>,
-  serverOptions: Partial<SafeJsonRpcServerOptions> = {},
+  serverOptions: SafeJsonRpcServerOptions = {},
   clientOptions: SafeJsonRpcClientOptions = {},
 ): Promise<JsonRpcTestContext<T>> {
   const server = createSafeJsonRpcServer<T>(handler, schema, serverOptions);

package/src/log/console.ts

@@ -7,7 +7,10 @@ import type { LogFn } from './log_fn.js';
  * which can be useful for controlling the format of the output or redirecting logs to a different destination.
  */
 class ConsoleLogger {
-  constructor(private prefix: string, private logger: (...args: any[]) => void = console.log) {}
+  constructor(
+    private prefix: string,
+    private logger: (...args: any[]) => void = console.log,
+  ) {}
 
   /**
    * Log messages with the specified prefix using the provided logger.

package/src/log/gcloud-logger-config.ts

@@ -13,10 +13,6 @@ const GOOGLE_CLOUD_TRACE_SAMPLED = 'logging.googleapis.com/trace_sampled';
  */
 export const GoogleCloudLoggerConfig = {
   messageKey: 'message',
-  // Same as pino.stdTimeFunctions.isoTime but uses "timestamp" key instead of "time"
-  timestamp(): string {
-    return `,"timestamp":"${new Date(Date.now()).toISOString()}"`;
-  },
   formatters: {
     log(object: Record<string, unknown>): Record<string, unknown> {
       // Add trace context attributes following Cloud Logging structured log format described
@@ -39,12 +35,12 @@ export const GoogleCloudLoggerConfig = {
       // Severity labels https://cloud.google.com/logging/docs/reference/v2/rest/v2/LogEntry#LogSeverity
       let severity: string;
 
-      switch (label as pino.Level | keyof typeof customLevels) {
+      switch (label as pino.Level | keyof CustomLevels) {
         case 'trace':
         case 'debug':
+        case 'verbose':
           severity = 'DEBUG';
           break;
-        case 'verbose':
         case 'info':
           severity = 'INFO';
           break;
@@ -68,4 +64,4 @@ export const GoogleCloudLoggerConfig = {
 } satisfies pino.LoggerOptions;
 
 // Define custom logging levels for pino. Duplicate from pino-logger.ts.
-const customLevels = { verbose: 25 };
+type CustomLevels = { verbose: 25 };

package/src/log/index.ts

@@ -1,6 +1,6 @@
 export * from './console.js';
-export * from './debug.js';
-export * from './pino-logger.js';
-export * from './log_history.js';
-export * from './log_fn.js';
 export * from './libp2p_logger.js';
+export * from './log_fn.js';
+export * from './noir_debug_log_util.js';
+export * from './pino-logger.js';
+export * from './log-levels.js';

package/src/log/libp2p_logger.ts

@@ -33,7 +33,7 @@ function createLibp2pLogger(component: string): Logger {
       return;
     }
 
-    log[level](replaceFormatting(message), ...args);
+    log[level](replaceFormatting(message), ...formatArgs(message, args));
   };
 
   // Default log level is trace as this is super super noisy
@@ -45,7 +45,7 @@ function createLibp2pLogger(component: string): Logger {
     enabled: log.isLevelEnabled('debug'),
     error(message: string, ...args: unknown[]) {
       // We write error outputs as debug as they are often expected, e.g. connection errors can happen in happy paths
-      logIfEnabled('debug', message, ...args);
+      logIfEnabled('debug', `error: ${message}`, ...args);
     },
 
     debug(message: string, ...args: unknown[]) {
@@ -65,3 +65,22 @@ function createLibp2pLogger(component: string): Logger {
     },
   });
 }
+
+function formatArgs(message: string, args: unknown[]) {
+  if (!args) {
+    return args;
+  }
+  return args.map(arg => {
+    if (
+      typeof arg === 'object' &&
+      arg &&
+      'err' in arg &&
+      arg.err instanceof Error &&
+      'type' in arg.err &&
+      arg.err.type === 'AbortError'
+    ) {
+      delete arg.err; // Remove the AbortError from the logs
+    }
+    return arg;
+  });
+}

package/src/log/log_fn.ts

@@ -1,5 +1,5 @@
 /** Structured log data to include with the message. */
-export type LogData = Record<string, string | number | bigint | boolean | { toString(): string } | undefined>;
+export type LogData = Record<string, string | number | bigint | boolean | { toString(): string } | undefined | null>;
 
 /** A callable logger instance. */
 export type LogFn = (msg: string, data?: unknown) => void;

package/src/log/noir_debug_log_util.ts

@@ -0,0 +1,21 @@
+interface Printable {
+  toString(): string;
+}
+
+/**
+ * Format a debug string filling in `'{0}'` entries with their
+ * corresponding values from the args array, amd `'{}'` with the whole array.
+ *
+ * @param formatStr - str of form `'this is a string with some entries like {0} and {1}'`
+ * @param args - array of fields to fill in the string format entries with
+ * @returns formatted string
+ */
+export function applyStringFormatting(formatStr: string, args: Printable[]): string {
+  return formatStr
+    .replace(/{(\d+)}/g, (match, index) => {
+      return typeof args[index] === 'undefined' ? match : args[index].toString();
+    })
+    .replace(/{}/g, (_match, _index) => {
+      return args.toString();
+    });
+}

package/src/log/pino-logger.ts

@@ -37,11 +37,14 @@ export function createLogger(module: string): Logger {
     debug: (msg: string, data?: unknown) => logFn('debug', msg, data),
     /** Log as trace. Use for when we want to denial-of-service any recipient of the logs. */
     trace: (msg: string, data?: unknown) => logFn('trace', msg, data),
+    /** Level of the logger */
     level: pinoLogger.level as LogLevel,
     /** Whether the given level is enabled for this logger. */
     isLevelEnabled: (level: LogLevel) => isLevelEnabled(pinoLogger, level),
     /** Module name for the logger. */
     module,
+    /** Creates another logger by extending this logger module name. */
+    createChild: (childModule: string) => createLogger(`${module}:${childModule}`),
   };
 }
 
@@ -99,11 +102,38 @@ const customLevels = { verbose: 25 };
 
 // Global pino options, tweaked for google cloud if running there.
 const useGcloudLogging = parseBooleanEnv(process.env['USE_GCLOUD_LOGGING' satisfies EnvVar]);
+
+const redactedPaths = [
+  'validatorPrivateKeys',
+  // for both the validator and the prover
+  'publisherPrivateKeys',
+  'peerIdPrivateKey',
+  // bot keys
+  'l1PrivateKey',
+  'senderPrivateKey',
+  // blob sink
+  'l1ConsensusHostApiKeys',
+  // sensitive options used in the CLI
+  'privateKey',
+  'mnemonic',
+  'l1Mnemonic',
+  'l1PrivateKey',
+];
+
 const pinoOpts: pino.LoggerOptions<keyof typeof customLevels> = {
   customLevels,
   messageKey: 'msg',
   useOnlyCustomLevels: false,
   level: logLevel,
+  redact: {
+    paths: [
+      ...redactedPaths,
+      ...redactedPaths.map(p => `config.${p}`),
+      ...redactedPaths.map(p => `cfg.${p}`),
+      ...redactedPaths.map(p => `options.${p}`),
+      ...redactedPaths.map(p => `opts.${p}`),
+    ],
+  },
   ...(useGcloudLogging ? GoogleCloudLoggerConfig : {}),
 };
 
@@ -221,7 +251,7 @@ export function registerLoggingStream(stream: Writable): void {
 }
 
 /** Log function that accepts an exception object */
-type ErrorLogFn = (msg: string, err?: Error | unknown, data?: LogData) => void;
+type ErrorLogFn = (msg: string, err?: unknown, data?: LogData) => void;
 
 /**
  * Logger that supports multiple severity levels.
@@ -230,6 +260,7 @@ export type Logger = { [K in LogLevel]: LogFn } & { /** Error log function */ er
   level: LogLevel;
   isLevelEnabled: (level: LogLevel) => boolean;
   module: string;
+  createChild: (childModule: string) => Logger;
 };
 
 /**
@@ -238,6 +269,6 @@ export type Logger = { [K in LogLevel]: LogFn } & { /** Error log function */ er
  * @param err - Error to log
  * @returns A string with both the log message and the error message.
  */
-function formatErr(msg: string, err?: Error | unknown): string {
+function formatErr(msg: string, err?: unknown): string {
   return err ? `${msg}: ${inspect(err)}` : msg;
 }

package/src/message/index.ts

@@ -31,7 +31,11 @@ interface TypedMessageLike {
 }
 
 export class TypedMessage<T, B> {
-  public constructor(public readonly msgType: T, public readonly header: MessageHeader, public readonly value: B) {}
+  public constructor(
+    public readonly msgType: T,
+    public readonly header: MessageHeader,
+    public readonly value: B,
+  ) {}
 
   static fromMessagePack<T, B>(data: TypedMessageLike): TypedMessage<T, B> {
     return new TypedMessage<T, B>(data['msgType'] as T, MessageHeader.fromMessagePack(data['header']), data['value']);

package/src/mutex/mutex_database.ts

@@ -3,10 +3,9 @@
  * Provides functionality for acquiring, extending, and releasing locks on resources to ensure exclusive access and prevent conflicts in concurrent applications.
  */
 export interface MutexDatabase {
-  // eslint-disable-next-line jsdoc/require-jsdoc
   acquireLock(name: string, timeout: number): Promise<boolean>;
-  // eslint-disable-next-line jsdoc/require-jsdoc
+
   extendLock(name: string, timeout: number): Promise<void>;
-  // eslint-disable-next-line jsdoc/require-jsdoc
+
   releaseLock(name: string): Promise<void>;
 }

package/src/number/index.ts

@@ -0,0 +1,14 @@
+/** Bounds the given value between the min and max provided (both inclusive) */
+export function boundInclusive(value: number, min: number, max: number): number {
+  if (min > max) {
+    throw new Error('Minimum bound cannot be greater than maximum bound');
+  }
+
+  if (value < min) {
+    return min;
+  } else if (value > max) {
+    return max;
+  } else {
+    return value;
+  }
+}

package/src/profiler/index.ts

@@ -0,0 +1 @@
+export * from './profiler.js';

package/src/profiler/profiler.ts

@@ -0,0 +1,125 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import { performance } from 'node:perf_hooks';
+
+interface Span {
+  label: string;
+  start: number;
+  dur: number;
+  count: number;
+  children: Span[];
+  parent: Span | undefined;
+}
+
+interface ProfileData {
+  spans: SerializedSpan[];
+  timestamp: string;
+  totalTime: number;
+}
+
+interface SerializedSpan {
+  label: string;
+  dur: number;
+  count: number;
+  children: SerializedSpan[];
+}
+
+const als = new AsyncLocalStorage<Span>();
+const roots: Span[] = [];
+
+function reset(): void {
+  roots.length = 0;
+}
+
+// Strip out circular references (parent) and unused fields (start) for JSON serialization
+function serializeSpans(spans: Span[]): SerializedSpan[] {
+  return spans.map(span => ({
+    label: span.label,
+    dur: span.dur,
+    count: span.count,
+    children: serializeSpans(span.children),
+  }));
+}
+
+let i = 0;
+function save(): void {
+  if (roots.length === 0) {
+    return;
+  }
+
+  // Find max single execution time across all spans (dur/count since dur is accumulated)
+  const findMaxSingleDuration = (spans: Span[]): number => {
+    let max = 0;
+    for (const span of spans) {
+      const singleDur = span.dur / span.count;
+      max = Math.max(max, singleDur);
+      if (span.children.length > 0) {
+        max = Math.max(max, findMaxSingleDuration(span.children));
+      }
+    }
+    return max;
+  };
+
+  const profileData: ProfileData = {
+    spans: serializeSpans(roots),
+    timestamp: new Date().toISOString(),
+    totalTime: findMaxSingleDuration(roots),
+  };
+
+  const profilePath = path.join(process.cwd(), `profile-${i++}.json`);
+  process.stdout.write(`Writing profile data to ${profilePath}\n`);
+  fs.writeFileSync(profilePath, JSON.stringify(profileData, null, 2));
+}
+
+// Hook into Jest to save after each test
+if (typeof afterEach === 'function') {
+  afterEach(() => {
+    save();
+    reset();
+  });
+}
+
+// Also save on process exit for non-Jest environments
+process.on('exit', () => {
+  save();
+});
+
+// Wrapper for async functions to maintain context properly
+async function runAsync<ReturnType>(label: string, fn: () => Promise<ReturnType>): Promise<ReturnType> {
+  const parent = als.getStore();
+
+  // Check if we already have a span with this label in the current context
+  let existingSpan: Span | undefined;
+  if (parent) {
+    existingSpan = parent.children.find(c => c.label === label);
+  } else {
+    existingSpan = roots.find(r => r.label === label);
+  }
+
+  let span: Span;
+  if (existingSpan) {
+    // Reuse existing span and increment count
+    span = existingSpan;
+    span.count++;
+  } else {
+    // Create new span
+    span = { label, start: performance.now(), dur: 0, count: 1, children: [], parent };
+    if (parent) {
+      parent.children.push(span);
+    } else {
+      roots.push(span);
+    }
+  }
+
+  const startTime = performance.now();
+  const result: ReturnType = await als.run(span, fn);
+  const elapsed = performance.now() - startTime;
+
+  // Add to total duration (for averaging)
+  span.dur += elapsed;
+
+  return result;
+}
+
+export const profiler = { reset, runAsync };