feat: Add initial tracing to SignatureController

packages/controller-utils/src/types.ts

@@ -127,3 +127,40 @@ export type NetworkNickname =
 // TODO: Move to @metamask/utils
 export type Partialize<Type, Key extends keyof Type> = Omit<Type, Key> &
   Partial<Pick<Type, Key>>;
+
+/** A context in which to execute a trace, in order to generate nested timings. */
+export type TraceContext = unknown;
+
+/** Request to trace an operation. */
+export type TraceRequest = {
+  /** Additional data to include in the trace. */
+  data?: Record<string, number | string | boolean>;
+
+  /** Name of the operation. */
+  name: string;
+
+  /**
+   * Unique identifier for the trace.
+   * Required if starting a trace and not providing a callback.
+   */
+  id?: string;
+
+  /** Trace context in which to execute the operation. */
+  parentContext?: TraceContext;
+
+  /** Additional tags to include in the trace to filter results. */
+  tags?: Record<string, number | string | boolean>;
+};
+
+/** Callback that traces the performance of an operation. */
+export type TraceCallback = <ReturnType>(
+  /** Request to trace the performance of an operation. */
+  request: TraceRequest,
+
+  /**
+   * Callback to trace.
+   * Thrown errors will not be caught, but the trace will still be recorded.
+   * @param context - The context in which the operation is running.
+   */
+  fn?: (context?: TraceContext) => ReturnType,
+) => Promise<ReturnType>;

packages/signature-controller/src/SignatureController.ts

@@ -10,6 +10,7 @@ import type {
 } from '@metamask/base-controller';
 import { BaseController } from '@metamask/base-controller';
 import { ApprovalType, ORIGIN_METAMASK } from '@metamask/controller-utils';
+import type { TraceCallback, TraceContext } from '@metamask/controller-utils';
 import type {
   KeyringControllerSignMessageAction,
   KeyringControllerSignPersonalMessageAction,
@@ -109,9 +110,9 @@ export type SignatureControllerMessenger = RestrictedControllerMessenger<
 >;
 
 export type SignatureControllerOptions = {
-  messenger: SignatureControllerMessenger;
-  isEthSignEnabled: () => boolean;
   getAllState: () => unknown;
+  getCurrentChainId: () => Hex;
+  messenger: SignatureControllerMessenger;
   securityProviderRequest?: (
     // TODO: Replace `any` with type
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -120,7 +121,7 @@ export type SignatureControllerOptions = {
     // TODO: Replace `any` with type
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
   ) => Promise<any>;
-  getCurrentChainId: () => Hex;
+  trace?: TraceCallback;
 };
 
 /**
@@ -141,6 +142,8 @@ export class SignatureController extends BaseController<
 
   #typedMessageManager: TypedMessageManager;
 
+  #trace: TraceCallback;
+
   /**
    * Construct a Sign controller.
    *
@@ -149,12 +152,14 @@ export class SignatureController extends BaseController<
    * @param options.getAllState - Callback to retrieve all user state.
    * @param options.securityProviderRequest - A function for verifying a message, whether it is malicious or not.
    * @param options.getCurrentChainId - A function for retrieving the current chainId.
+   * @param options.trace - Callback to generate trace information.
    */
   constructor({
-    messenger,
     getAllState,
-    securityProviderRequest,
     getCurrentChainId,
+    messenger,
+    securityProviderRequest,
+    trace,
   }: SignatureControllerOptions) {
     super({
       name: controllerName,
@@ -178,6 +183,7 @@ export class SignatureController extends BaseController<
       undefined,
       getCurrentChainId,
     );
+    this.#trace = trace ?? (((_request, fn) => fn?.()) as TraceCallback);
 
     this.#handleMessageManagerEvents(
       this.#personalMessageManager,
@@ -276,11 +282,14 @@ export class SignatureController extends BaseController<
    *
    * @param messageParams - The params of the message to sign & return to the Dapp.
    * @param req - The original request, containing the origin.
+   * @param options - An options bag for the method.
+   * @param options.traceContext - The parent context for any new traces.
    * @returns Promise resolving to the raw data of the signature request.
    */
   async newUnsignedPersonalMessage(
     messageParams: PersonalMessageParams,
     req: OriginalRequest,
+    options: { traceContext?: TraceContext } = {},
   ): Promise<string> {
     return this.#newUnsignedAbstractMessage(
       this.#personalMessageManager,
@@ -292,6 +301,9 @@ export class SignatureController extends BaseController<
       this.#signPersonalMessage.bind(this),
       messageParams,
       req,
+      undefined,
+      undefined,
+      options.traceContext,
     );
   }
 
@@ -303,13 +315,16 @@ export class SignatureController extends BaseController<
    * @param version - The version indicating the format of the typed data.
    * @param signingOpts - An options bag for signing.
    * @param signingOpts.parseJsonData - Whether to parse the JSON before signing.
+   * @param options - An options bag for the method.
+   * @param options.traceContext - The parent context for any new traces.
    * @returns Promise resolving to the raw data of the signature request.
    */
   async newUnsignedTypedMessage(
     messageParams: TypedMessageParams,
     req: OriginalRequest,
     version: string,
     signingOpts: TypedMessageSigningOptions,
+    options: { traceContext?: TraceContext } = {},
   ): Promise<string> {
     const signTypeForLogger = this.#getSignTypeForLogger(version);
     return this.#newUnsignedAbstractMessage(
@@ -324,6 +339,7 @@ export class SignatureController extends BaseController<
       req,
       version,
       signingOpts,
+      options.traceContext,
     );
   }
 
@@ -397,6 +413,7 @@ export class SignatureController extends BaseController<
     req: OriginalRequest,
     version?: string,
     signingOpts?: SO,
+    traceContext?: TraceContext,
   ) {
     let resultCallbacks: AcceptResultCallbacks | undefined;
     try {
@@ -425,9 +442,13 @@ export class SignatureController extends BaseController<
           messageParamsWithId,
         );
 
-        const acceptResult = await this.#requestApproval(
-          messageParamsWithId,
-          approvalType,
+        const acceptResult = await this.#trace(
+          { name: 'Await Approval', parentContext: traceContext },
+          (context) =>
+            this.#requestApproval(messageParamsWithId, approvalType, {
+              traceContext: context,
+              actionId: req?.id?.toString(),
+            }),
         );
 
         resultCallbacks = acceptResult.resultCallbacks;
@@ -445,7 +466,10 @@ export class SignatureController extends BaseController<
 
       // TODO: Either fix this lint violation or explain why it's necessary to ignore.
       // eslint-disable-next-line @typescript-eslint/await-thenable
-      await signMessage(messageParamsWithId, signingOpts);
+
+      await this.#trace({ name: 'Sign', parentContext: traceContext }, () =>
+        signMessage(messageParamsWithId, signingOpts),
+      );
 
       const signatureResult = await signaturePromise;
 
@@ -807,13 +831,24 @@ export class SignatureController extends BaseController<
   async #requestApproval(
     msgParams: AbstractMessageParamsMetamask,
     type: ApprovalType,
+    {
+      traceContext,
+      actionId,
+    }: { traceContext?: TraceContext; actionId?: string },
   ): Promise<AddResult> {
     const id = msgParams.metamaskId as string;
     const origin = msgParams.origin || ORIGIN_METAMASK;
 
     // We are explicitly cloning the message params here to prevent the mutation errors on development mode
     // Because sending it through the messaging system will make the object read only
     const clonedMsgParams = cloneDeep(msgParams);
+
+    await this.#trace({
+      name: 'Notification Display',
+      id: actionId,
+      parentContext: traceContext,
+    });
+
     return (await this.messagingSystem.call(
       'ApprovalController:addRequest',
       {

packages/transaction-controller/src/TransactionController.ts

@@ -21,6 +21,7 @@ import {
   convertHexToDecimal,
   isInfuraNetworkType,
 } from '@metamask/controller-utils';
+import type { TraceCallback, TraceContext } from '@metamask/controller-utils';
 import EthQuery from '@metamask/eth-query';
 import type {
   FetchGasFeeEstimateOptions,
@@ -78,8 +79,6 @@ import type {
   SimulationData,
   GasFeeEstimates,
   GasFeeFlowResponse,
-  TraceCallback,
-  TraceContext,
   GasPriceValue,
   FeeMarketEIP1559Values,
 } from './types';

packages/transaction-controller/src/types.ts

@@ -1300,43 +1300,6 @@ export type SimulationData = {
   tokenBalanceChanges: SimulationTokenBalanceChange[];
 };
 
-/** A context in which to execute a trace, in order to generate nested timings. */
-export type TraceContext = unknown;
-
-/** Request to trace an operation. */
-export type TraceRequest = {
-  /** Additional data to include in the trace. */
-  data?: Record<string, number | string | boolean>;
-
-  /** Name of the operation. */
-  name: string;
-
-  /**
-   * Unique identifier for the trace.
-   * Required if starting a trace and not providing a callback.
-   */
-  id?: string;
-
-  /** Trace context in which to execute the operation. */
-  parentContext?: TraceContext;
-
-  /** Additional tags to include in the trace to filter results. */
-  tags?: Record<string, number | string | boolean>;
-};
-
-/** Callback that traces the performance of an operation. */
-export type TraceCallback = <ReturnType>(
-  /** Request to trace the performance of an operation. */
-  request: TraceRequest,
-
-  /**
-   * Callback to trace.
-   * Thrown errors will not be caught, but the trace will still be recorded.
-   * @param context - The context in which the operation is running.
-   */
-  fn?: (context?: TraceContext) => ReturnType,
-) => Promise<ReturnType>;
-
 /** Gas fee properties for a legacy transaction. */
 export type GasPriceValue = {
   /** Price per gas for legacy transactions. */