From b903e6cff06bc3d2558bec34e273c7dff633ef50 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Tue, 28 Jul 2026 21:46:05 +0100 Subject: [PATCH] feat(assets-controller): enhance Sentry tracing for asset fetch and update pipelines - Nested Sentry spans for asset fetching under a single root span per pipeline, reducing transaction count. - Updated `AssetsFullFetch`, `AssetsControllerFirstInitFetch`, and related spans to be subspans of `AssetsFetchPipeline` and `AssetsBackgroundFetch`. - Adjusted `AssetsFullFetch.duration_ms` to measure middleware execution only. - Implemented best-effort tracing for rejected promises in asset updates, ensuring underlying work still executes. - Added tests to verify new tracing behavior and ensure correct span nesting. This change improves performance monitoring and error handling in asset management workflows. --- packages/assets-controller/CHANGELOG.md | 7 + .../src/AssetsController.test.ts | 249 ++++++++- .../assets-controller/src/AssetsController.ts | 379 ++++++++----- .../src/selectors/balance.test.ts | 52 ++ .../src/selectors/balance.ts | 80 ++- packages/assets-controller/src/traced.test.ts | 500 ++++++++++++++++++ packages/assets-controller/src/traced.ts | 293 ++++++++++ 7 files changed, 1425 insertions(+), 135 deletions(-) create mode 100644 packages/assets-controller/src/traced.test.ts create mode 100644 packages/assets-controller/src/traced.ts diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index afe58b92d50..a1479de23f4 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -13,12 +13,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- Nest the Sentry spans emitted while fetching assets under one root span per pipeline, so that a fetch costs a single Sentry transaction instead of one per measurement ([#9672](https://github.com/MetaMask/core/pull/9672)) + - `AssetsFullFetch`, `AssetsControllerFirstInitFetch`, `AssetsDataSourceTiming` and `AssetsDataSourceError` are now recorded as subspans of `AssetsFetchPipeline` (fast lane) or `AssetsBackgroundFetch` (background lane), and `AssetsUpdatePipeline` as a subspan of `AssetsUpdateEnrichment`. + - Pipeline spans are emitted only on the unlock (first-init) fetch of a session. Later polls, force updates and subscription-driven enrichment no longer emit them. `AssetsFullFetch` previously fired on every force update. + - `AssetsFullFetch.duration_ms` now measures middleware execution only, and no longer includes building the data request or committing the result to state. + - Dashboard-facing spans copy numeric span data into tags and backdate `startTime`, so Sentry records them as measurements that Spans widgets charting `p95(duration_ms)` can read. + - `calculateBalanceForAllWallets` emits a single `AggregatedBalanceSelector` subspan under an `AggregatedBalance` parent, instead of one root span per account group. - Bump `@metamask/keyring-api` from `^23.5.0` to `^23.7.0` ([#9676](https://github.com/MetaMask/core/pull/9676)) - Bump `@metamask/keyring-internal-api` from `^11.0.1` to `^11.0.2` ([#9676](https://github.com/MetaMask/core/pull/9676)) - Bump `@metamask/keyring-snap-client` from `^9.2.0` to `^9.2.1` ([#9676](https://github.com/MetaMask/core/pull/9676)) ### Fixed +- A rejected `trace` promise can no longer fail a full fetch or `handleAssetsUpdate` enrichment; tracing is now best-effort throughout, and the underlying work still runs exactly once ([#9672](https://github.com/MetaMask/core/pull/9672)) - `SnapDataSource` now delivers snap-sourced balance updates directly to `AssetsController` via a constructor-supplied `onAssetsUpdate` callback instead of fanning out to `activeSubscriptions`, so updates (e.g. Tron energy/bandwidth) are no longer dropped when no active subscription is tracked for the chain in the SnapDataSource ([#9656](https://github.com/MetaMask/core/pull/9656)) ## [11.2.1] diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 3afadff79af..65ec2442016 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -1395,6 +1395,98 @@ describe('AssetsController', () => { }); describe('handleAssetsUpdate', () => { + it('nests the update pipeline record under an AssetsUpdateEnrichment root', async () => { + const enrichmentContext = { id: 'assets-update-enrichment' }; + const traceMock = jest + .fn() + .mockImplementation( + async ( + request: TraceRequest, + fn?: (context?: unknown) => unknown, + ) => { + return fn?.( + request.parentContext === undefined + ? enrichmentContext + : undefined, + ); + }, + ); + const trace = traceMock as unknown as TraceCallback; + + await withController( + { controllerOptions: { trace } }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'merge', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [MOCK_ASSET_ID]: { amount: '1' } }, + }, + }, + 'test-source', + ); + + const requests = traceMock.mock.calls.map( + ([request]) => request as TraceRequest, + ); + const byName = (name: string): TraceRequest[] => + requests.filter((request) => request.name === name); + + expect(byName('AssetsUpdateEnrichment')).toMatchObject([ + { parentContext: undefined }, + ]); + expect(byName('AssetsUpdatePipeline')).toMatchObject([ + { + parentContext: enrichmentContext, + data: expect.objectContaining({ + source: 'test-source', + duration_ms: expect.any(Number), + has_balance: true, + }), + startTime: expect.any(Number), + }, + ]); + }, + ); + }); + + it('does not fail when the tracer rejects after enrichment completes', async () => { + const trace = jest + .fn() + .mockImplementation( + async ( + _request: TraceRequest, + fn?: (context?: unknown) => unknown, + ) => { + if (fn) { + await fn({ id: 'parent' }); + throw new Error('telemetry failed'); + } + return undefined; + }, + ) as unknown as TraceCallback; + + await withController( + { controllerOptions: { trace } }, + async ({ controller }) => { + expect( + await controller.handleAssetsUpdate( + { + updateMode: 'merge', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [MOCK_ASSET_ID]: { amount: '1' } }, + }, + }, + 'test-source', + ), + ).toBeUndefined(); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toStrictEqual({ amount: '1' }); + }, + ); + }); + it('preserves existing rich metadata when the API response has empty symbol and name', async () => { const richMetadata: FungibleAssetMetadata = { type: 'erc20', @@ -2519,7 +2611,13 @@ describe('AssetsController', () => { duration_ms: expect.any(Number), chain_ids: expect.any(String), }), - tags: { controller: 'AssetsController' }, + // Numeric data is mirrored into tags so the Sentry adapter records + // it as a measurement, and the span is backdated to match. + tags: expect.objectContaining({ + controller: 'AssetsController', + duration_ms: expect.any(Number), + }), + startTime: expect.any(Number), }); const { duration_ms: durationMs, @@ -2533,6 +2631,112 @@ describe('AssetsController', () => { ); }); + it('nests fetch spans under a single AssetsFetchPipeline root on unlock', async () => { + const pipelineContext = { id: 'assets-fetch-pipeline' }; + const traceMock = jest + .fn() + .mockImplementation( + async ( + request: TraceRequest, + fn?: (context?: unknown) => unknown, + ) => { + // Only the root span (no parent of its own) hands out a context. + return fn?.( + request.parentContext === undefined ? pipelineContext : undefined, + ); + }, + ); + const trace = traceMock as unknown as TraceCallback; + + await withController( + { + clientControllerState: { isUiOpen: true }, + controllerOptions: { trace }, + }, + async ({ messenger }) => { + ( + messenger as unknown as { + publish: (topic: string, payload?: unknown) => void; + } + ).publish('ClientController:stateChange', { isUiOpen: true }); + messenger.publish('KeyringController:unlock'); + await new Promise((resolve) => setTimeout(resolve, 100)); + + const requests = traceMock.mock.calls.map( + ([request]) => request as TraceRequest, + ); + const byName = (name: string): TraceRequest[] => + requests.filter((request) => request.name === name); + + expect(byName('AssetsFetchPipeline')).toHaveLength(1); + expect(byName('AssetsFetchPipeline')[0]).toMatchObject({ + parentContext: undefined, + tags: { controller: 'AssetsController' }, + }); + + for (const name of [ + 'AssetsFullFetch', + 'AssetsControllerFirstInitFetch', + 'AssetsDataSourceTiming', + ]) { + const nested = byName(name); + expect(nested.length).toBeGreaterThan(0); + for (const request of nested) { + expect(request.parentContext).toBe(pipelineContext); + } + } + }, + ); + }); + + it('does not fail a force update when the tracer rejects after the work completes', async () => { + const trace = jest + .fn() + .mockImplementation( + async ( + _request: TraceRequest, + fn?: (context?: unknown) => unknown, + ) => { + if (fn) { + await fn({ id: 'parent' }); + // Simulate a Sentry adapter failing once the span closes. + throw new Error('telemetry failed'); + } + return undefined; + }, + ) as unknown as TraceCallback; + + await withController( + { controllerOptions: { trace } }, + async ({ controller }) => { + expect( + await controller.getAssets([createMockInternalAccount()], { + forceUpdate: true, + }), + ).toBeDefined(); + }, + ); + }); + + it('still runs a force update when the tracer rejects before invoking the work', async () => { + const trace = jest + .fn() + .mockRejectedValue( + new Error('telemetry failed'), + ) as unknown as TraceCallback; + + await withController( + { controllerOptions: { trace } }, + async ({ controller }) => { + expect( + await controller.getAssets([createMockInternalAccount()], { + forceUpdate: true, + }), + ).toBeDefined(); + }, + ); + }); + it('replaces pre-lock balances on unlock via merge with covered-chain replacement', async () => { const fetchV5MultiAccountBalances = jest.fn().mockResolvedValue({ balances: [ @@ -2627,6 +2831,49 @@ describe('AssetsController', () => { }, ); }); + + it('stops emitting pipeline spans once the first-init fetch has reported', async () => { + const traceMock = jest + .fn() + .mockImplementation( + async (_request: TraceRequest, fn?: (context?: unknown) => unknown) => + fn?.(), + ); + const trace = traceMock as unknown as TraceCallback; + + await withController( + { + clientControllerState: { isUiOpen: true }, + controllerOptions: { trace }, + }, + async ({ controller, messenger }) => { + ( + messenger as unknown as { + publish: (topic: string, payload?: unknown) => void; + } + ).publish('ClientController:stateChange', { isUiOpen: true }); + messenger.publish('KeyringController:unlock'); + await new Promise((resolve) => setTimeout(resolve, 100)); + + const countSpans = (): Record => { + const counts: Record = {}; + for (const [request] of traceMock.mock.calls) { + const { name } = request as TraceRequest; + counts[name] = (counts[name] ?? 0) + 1; + } + return counts; + }; + const before = countSpans(); + + // Steady-state polling must not keep paying for spans. + await controller.getAssets([createMockInternalAccount()], { + forceUpdate: true, + }); + + expect(countSpans()).toStrictEqual(before); + }, + ); + }); }); describe('subscribeAssetsPrice', () => { diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index e9d7dfa01d9..f4f1c4e2764 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -115,6 +115,8 @@ import { import { RpcFallbackMiddleware } from './middlewares/RpcFallbackMiddleware.js'; import type { Assets3346MigrationState } from './migrations/healAssetsInfoMetadata.js'; import { tempHealAssetsInfoMetadata } from './migrations/healAssetsInfoMetadata.js'; +import type { SpanHandle, SummaryRecord } from './traced.js'; +import { emitTrace, registerTracer, traced } from './traced.js'; import type { AccountId, AssetPreferences, @@ -217,6 +219,159 @@ const TRACE_UPDATE_PIPELINE = 'AssetsUpdatePipeline'; const TRACE_SUBSCRIPTION_ERROR = 'AssetsSubscriptionError'; const TRACE_STATE_SIZE = 'AssetsStateSize'; +/** Tags applied to every span this controller emits. */ +const TRACE_TAGS = { controller: 'AssetsController' }; + +/** + * Root span names for each pipeline lane. The per-source timings and the + * dashboard summary records for a lane nest underneath its root, so a fetch + * costs Sentry one transaction instead of one per measurement. + */ +const LANE_SPAN_NAMES = { + fast: 'AssetsFetchPipeline', + background: 'AssetsBackgroundFetch', + update: 'AssetsUpdateEnrichment', +} as const; + +/** + * Arguments for one `#executeMiddlewares` run. The lane discriminates which + * extra values the dashboard records for that pipeline need. + */ +type ExecuteMiddlewaresOptions = { + /** Data sources or middlewares with `getName()` and `assetsMiddleware`, executed in order. */ + sources: AssetsDataSource[]; + /** The data request. */ + request: DataRequest; + /** Initial response, when enriching data that has already been fetched. */ + initialResponse?: DataResponse; +} & ( + | { + lane: 'fast'; + /** Every requested chain, including those left to the background lane. */ + chainIds: ChainId[]; + accountCount: number; + basicFunctionality: boolean; + } + | { lane: 'background'; accountCount: number } + | { lane: 'update'; sourceId: string } +); + +type ExecuteMiddlewaresResult = { + response: DataResponse; + durationByDataSource: Record; +}; + +/** + * Count how many balances a response carries across all accounts. + * + * @param response - The response to measure. + * @returns The total number of balance entries. + */ +function countBalanceEntries(response: DataResponse): number { + return Object.values(response.assetsBalance ?? {}).reduce( + (total, account) => total + Object.keys(account).length, + 0, + ); +} + +/** + * Build the data attached to a lane's root span, from the call's arguments. + * + * @param options - Arguments the pipeline was invoked with. + * @returns Span data for the lane. + */ +function buildLaneSpanData( + options: ExecuteMiddlewaresOptions, +): Record { + switch (options.lane) { + case 'fast': + return { + chain_count: options.chainIds.length, + account_count: options.accountCount, + basic_functionality: options.basicFunctionality, + }; + case 'background': + return { + chain_count: options.request.chainIds.length, + account_count: options.accountCount, + }; + default: + return { + source: options.sourceId, + has_balance: Boolean(options.initialResponse?.assetsBalance), + has_price: Boolean(options.initialResponse?.assetsPrice), + balance_account_count: Object.keys( + options.initialResponse?.assetsBalance ?? {}, + ).length, + }; + } +} + +/** + * Build the Assets Health dashboard records for a completed pipeline run. + * + * These are derived from the result, so `traced` emits them while the lane's + * root span is still open — a child attached to a finished Sentry transaction + * is orphaned and dropped. + * + * @param result - What the pipeline returned. + * @param elapsedMs - How long the pipeline took. + * @param options - Arguments the pipeline was invoked with. + * @returns Records to emit under the lane's root span. + */ +function buildLaneSummaryRecords( + result: ExecuteMiddlewaresResult, + elapsedMs: number, + options: ExecuteMiddlewaresOptions, +): SummaryRecord[] { + switch (options.lane) { + case 'fast': + return [ + { + name: TRACE_FULL_FETCH, + data: { + duration_ms: elapsedMs, + chain_count: options.chainIds.length, + account_count: options.accountCount, + basic_functionality: options.basicFunctionality, + asset_count: countBalanceEntries(result.response), + price_count: Object.keys(result.response.assetsPrice ?? {}).length, + ...result.durationByDataSource, + }, + tags: TRACE_TAGS, + }, + { + name: TRACE_FIRST_INIT_FETCH, + data: { + duration_ms: elapsedMs, + chain_ids: JSON.stringify(options.chainIds), + ...result.durationByDataSource, + }, + tags: TRACE_TAGS, + }, + ]; + case 'background': + return []; + default: + return [ + { + name: TRACE_UPDATE_PIPELINE, + data: { + source: options.sourceId, + duration_ms: elapsedMs, + has_balance: Boolean(options.initialResponse?.assetsBalance), + has_price: Boolean(options.initialResponse?.assetsPrice), + has_metadata: Boolean(result.response.assetsInfo), + balance_account_count: Object.keys( + options.initialResponse?.assetsBalance ?? {}, + ).length, + }, + tags: TRACE_TAGS, + }, + ]; + } +} + const log = createModuleLogger(projectLogger, CONTROLLER_NAME); // ============================================================================ @@ -664,28 +819,6 @@ export class AssetsController extends BaseController< /** Whether we have already reported state size for this session (reset on #stop). */ #stateSizeReported = false; - /** - * Fire-and-forget trace helper. Swallows errors so telemetry never breaks the controller. - * - * @param name - Trace / span name visible in Sentry. - * @param data - Key-value pairs attached as span data. - * @param tags - Key-value pairs used for Sentry filtering. - */ - #emitTrace( - name: string, - data: Record, - tags: Record = { - controller: 'AssetsController', - }, - ): void { - if (!this.#trace) { - return; - } - this.#trace({ name, data, tags }, () => undefined)?.catch(() => { - // Telemetry failure must not break. - }); - } - /** * Emit a state-size trace once on app start (first state update after unlock). */ @@ -727,14 +860,18 @@ export class AssetsController extends BaseController< } } - this.#emitTrace(TRACE_STATE_SIZE, { - balance_entries: balanceEntries, - balance_accounts: Object.keys(balances).length, - unique_asset_count: uniqueAssets.size, - network_count: uniqueNetworks.size, - metadata_entries: Object.keys(assetsInfo).length, - price_entries: Object.keys(assetsPrice).length, - custom_asset_entries: customAssetEntries, + emitTrace(this.#trace, undefined, { + name: TRACE_STATE_SIZE, + data: { + balance_entries: balanceEntries, + balance_accounts: Object.keys(balances).length, + unique_asset_count: uniqueAssets.size, + network_count: uniqueNetworks.size, + metadata_entries: Object.keys(assetsInfo).length, + price_entries: Object.keys(assetsPrice).length, + custom_asset_entries: customAssetEntries, + }, + tags: TRACE_TAGS, }); } @@ -872,6 +1009,9 @@ export class AssetsController extends BaseController< this.#isBasicFunctionality = isBasicFunctionality ?? ((): boolean => true); this.#defaultUpdateInterval = defaultUpdateInterval; this.#trace = trace; + // `@traced` methods read the tracer from here: a decorator wrapper cannot + // reach `this.#trace`, since private names are scoped to the class body. + registerTracer(this, trace); this.#captureException = captureException; this.#queryApiClient = queryApiClient; const rpcConfig = rpcDataSourceConfig ?? {}; @@ -1498,19 +1638,27 @@ export class AssetsController extends BaseController< * Execute middlewares with request/response context. * Returns response and exclusive duration per source (sum ≈ wall time). * - * @param sources - Data sources or middlewares with getName() and assetsMiddleware (executed in order). - * @param request - The data request. - * @param initialResponse - Optional initial response (for enriching existing data). + * This is the single traced choke point for all three pipelines: `@traced` + * opens the lane's root span, and everything this method emits through + * `span` nests underneath it. The handle is required rather than optional so + * that every call site has to state whether it wants tracing — an accidental + * new polling caller cannot silently start flooding Sentry. + * + * @param span - Handle of the caller's span; see `./traced.ts` for the protocol. + * @param options - The lane, sources, request and any initial response. * @returns Response and durationByDataSource (exclusive ms per source name). */ + @traced<[ExecuteMiddlewaresOptions], ExecuteMiddlewaresResult>({ + name: (options) => LANE_SPAN_NAMES[options.lane], + tags: TRACE_TAGS, + data: buildLaneSpanData, + summary: buildLaneSummaryRecords, + }) async #executeMiddlewares( - sources: AssetsDataSource[], - request: DataRequest, - initialResponse: DataResponse = {}, - ): Promise<{ - response: DataResponse; - durationByDataSource: Record; - }> { + span: SpanHandle, + options: ExecuteMiddlewaresOptions, + ): Promise { + const { sources, request, initialResponse = {} } = options; const names = sources.map((source) => source.getName()); const middlewares = sources.map((source) => source.assetsMiddleware); const inclusive: number[] = []; @@ -1576,15 +1724,21 @@ export class AssetsController extends BaseController< } } - // Emit per-source timing traces for the Assets Health dashboard + // Per-source timings for the Assets Health dashboard, nested under this + // lane's root span. for (const [sourceName, durationMs] of Object.entries( durationByDataSource, )) { - this.#emitTrace(TRACE_DATA_SOURCE_TIMING, { - source: sourceName, - duration_ms: durationMs, - chain_count: request.chainIds.length, - account_count: request.accountsWithSupportedChains.length, + emitTrace(this.#trace, span, { + name: TRACE_DATA_SOURCE_TIMING, + data: { + source: sourceName, + duration_ms: durationMs, + chain_count: request.chainIds.length, + account_count: request.accountsWithSupportedChains.length, + }, + // String tag so Spans widgets can group by source. + tags: { ...TRACE_TAGS, source: sourceName }, }); } @@ -1601,19 +1755,19 @@ export class AssetsController extends BaseController< } catch { // Never let telemetry throw. } - this.#emitTrace( - TRACE_DATA_SOURCE_ERROR, - { + emitTrace(this.#trace, span, { + name: TRACE_DATA_SOURCE_ERROR, + data: { failed_sources: failedSources, error_count: middlewareErrors.length, chain_count: request.chainIds.length, }, - { - controller: 'AssetsController', + tags: { + ...TRACE_TAGS, severity: 'error', error_type: assetsError.name, }, - ); + }); } return { response: result.response, durationByDataSource }; @@ -1651,7 +1805,11 @@ export class AssetsController extends BaseController< } if (options?.forceUpdate) { - const startTime = performance.now(); + // Pipeline spans are emitted only on the unlock/first-init fetch. Later + // polls and force updates reuse the same code path with tracing off, so + // that steady-state polling cannot flood Sentry. + const span: SpanHandle = { enabled: !this.#firstInitFetchReported }; + const request = this.#buildDataRequest(accounts, chainIds, { assetTypes, dataTypes, @@ -1685,10 +1843,17 @@ export class AssetsController extends BaseController< ]), ] : [this.#stakedBalanceDataSource, this.#detectionMiddleware]; - const { response, durationByDataSource } = await this.#executeMiddlewares( - fastSources, + const { response } = await this.#executeMiddlewares(span, { + lane: 'fast', + sources: fastSources, request, - ); + chainIds, + accountCount: accounts.length, + basicFunctionality: this.#isBasicFunctionality(), + }); + // Mark after the fast lane has reported, so a concurrent path cannot + // skip the first-init records. + this.#firstInitFetchReported = true; await this.#updateState({ ...response, updateMode: 'merge', @@ -1699,6 +1864,8 @@ export class AssetsController extends BaseController< // commits to state. Their balances are merged together before detection. // Token + price enrichment matches the pre-split behavior: only when basic // functionality is on (RPC-only mode must not call token/price APIs). + // It gets its own root span, because by the time it finishes the fast + // lane's span has already closed. const slowPipelineChainIds = this.#getSlowPipelineChainIds( chainIds, response, @@ -1709,57 +1876,31 @@ export class AssetsController extends BaseController< ? [this.#snapDataSource, this.#rpcDataSource] : [this.#rpcDataSource]; - const slowRequest = { ...request, chainIds: slowPipelineChainIds }; - this.#executeMiddlewares( - [ - createParallelBalanceMiddleware(slowSources), - this.#detectionMiddleware, - ...(this.#isBasicFunctionality() - ? [ - createParallelMiddleware([ - this.#tokenDataSource, - this.#priceDataSource, - ]), - ] - : []), - ], - slowRequest, + { enabled: span.enabled }, + { + lane: 'background', + sources: [ + createParallelBalanceMiddleware(slowSources), + this.#detectionMiddleware, + ...(this.#isBasicFunctionality() + ? [ + createParallelMiddleware([ + this.#tokenDataSource, + this.#priceDataSource, + ]), + ] + : []), + ], + request: { ...request, chainIds: slowPipelineChainIds }, + accountCount: accounts.length, + }, ) .then(({ response: slowResponse }) => this.#updateState({ ...slowResponse, updateMode: 'merge' }), ) .catch((error) => log('Background pipeline failed', { error })); } - - const durationMs = performance.now() - startTime; - - // Emit trace for every full fetch (Assets Health dashboard) - this.#emitTrace(TRACE_FULL_FETCH, { - duration_ms: durationMs, - chain_count: chainIds.length, - account_count: accounts.length, - basic_functionality: this.#isBasicFunctionality(), - asset_count: response.assetsBalance - ? Object.values(response.assetsBalance).reduce( - (sum, acct) => sum + Object.keys(acct).length, - 0, - ) - : 0, - price_count: response.assetsPrice - ? Object.keys(response.assetsPrice).length - : 0, - ...durationByDataSource, - }); - - if (!this.#firstInitFetchReported) { - this.#firstInitFetchReported = true; - this.#emitTrace(TRACE_FIRST_INIT_FETCH, { - duration_ms: durationMs, - chain_ids: JSON.stringify(chainIds), - ...durationByDataSource, - }); - } } const result = this.#getAssetsFromState(accounts, chainIds, assetTypes); @@ -3285,9 +3426,13 @@ export class AssetsController extends BaseController< `[AssetsController] Failed to subscribe to '${sourceId}':`, error, ); - this.#emitTrace(TRACE_SUBSCRIPTION_ERROR, { - source: sourceId, - error_message: String(error), + emitTrace(this.#trace, undefined, { + name: TRACE_SUBSCRIPTION_ERROR, + data: { + source: sourceId, + error_message: String(error), + }, + tags: TRACE_TAGS, }); }); @@ -3678,14 +3823,16 @@ export class AssetsController extends BaseController< sourceId: string, request?: DataRequest, ): Promise { - const updateStart = performance.now(); - log('Assets updated from data source', { sourceId, hasBalance: Boolean(response.assetsBalance), hasPrice: Boolean(response.assetsPrice), }); + // Enrichment spans only before the unlock/first-init fetch has reported; + // steady-state subscription updates would otherwise flood Sentry. + const span: SpanHandle = { enabled: !this.#firstInitFetchReported }; + const resolvedRequest: DataRequest = request ?? { accountsWithSupportedChains: [], chainIds: [], @@ -3727,26 +3874,20 @@ export class AssetsController extends BaseController< } const { response: enrichedResponse } = await this.#executeMiddlewares( - enrichmentSources, - pipelineRequest, - response, + span, + { + lane: 'update', + sources: enrichmentSources, + request: pipelineRequest, + initialResponse: response, + sourceId, + }, ); await this.#updateState({ ...enrichedResponse, replaceCoveredChainBalances: response.replaceCoveredChainBalances, }); - - this.#emitTrace(TRACE_UPDATE_PIPELINE, { - source: sourceId, - duration_ms: performance.now() - updateStart, - has_balance: Boolean(response.assetsBalance), - has_price: Boolean(response.assetsPrice), - has_metadata: Boolean(enrichedResponse.assetsInfo), - balance_account_count: response.assetsBalance - ? Object.keys(response.assetsBalance).length - : 0, - }); } // ============================================================================ diff --git a/packages/assets-controller/src/selectors/balance.test.ts b/packages/assets-controller/src/selectors/balance.test.ts index facc9a9af5b..fd779423eea 100644 --- a/packages/assets-controller/src/selectors/balance.test.ts +++ b/packages/assets-controller/src/selectors/balance.test.ts @@ -724,6 +724,58 @@ describe('wallet-balance selectors', () => { expect(result.totalBalanceInUserCurrency).toBe(0); }); + + it('emits one AggregatedBalanceSelector subspan for the whole walk', () => { + // Forwarding `trace` per group used to open one root span per account + // group, which hit Sentry's rate limits. + const parentSpan = { id: 'aggregated-balance' }; + const trace = jest + .fn() + .mockImplementation( + async ( + request: { parentContext?: unknown }, + fn?: (context?: unknown) => unknown, + ) => + fn?.(request.parentContext === undefined ? parentSpan : undefined), + ); + + calculateBalanceForAllWallets( + buildState(), + accountTreeState, + undefined, + trace, + ); + + expect(trace).toHaveBeenCalledTimes(2); + expect(trace).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ + name: 'AggregatedBalance', + data: expect.objectContaining({ + duration_ms: expect.any(Number), + wallet_count: 2, + group_count: 3, + }), + tags: expect.objectContaining({ + controller: 'AssetsController', + duration_ms: expect.any(Number), + }), + startTime: expect.any(Number), + }), + expect.any(Function), + ); + expect(trace).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + name: 'AggregatedBalanceSelector', + parentContext: parentSpan, + data: expect.objectContaining({ duration_ms: expect.any(Number) }), + tags: expect.objectContaining({ duration_ms: expect.any(Number) }), + startTime: expect.any(Number), + }), + expect.any(Function), + ); + }); }); describe('calculateBalanceChangeForAccountGroup', () => { diff --git a/packages/assets-controller/src/selectors/balance.ts b/packages/assets-controller/src/selectors/balance.ts index c993af84e4e..62c6769588d 100644 --- a/packages/assets-controller/src/selectors/balance.ts +++ b/packages/assets-controller/src/selectors/balance.ts @@ -12,6 +12,7 @@ import { import BigNumberJS from 'bignumber.js'; import type { AssetsControllerState } from '../AssetsController.js'; +import { emitTrace } from '../traced.js'; import type { AccountId, AssetBalance, @@ -23,8 +24,55 @@ import type { // ============================================================================ // TRACE NAMES — used in Sentry spans (search these strings in Discover) // ============================================================================ +/** Parent span that nests {@link TRACE_AGGREGATED_BALANCE_SELECTOR}. */ +const TRACE_AGGREGATED_BALANCE = 'AggregatedBalance'; const TRACE_AGGREGATED_BALANCE_SELECTOR = 'AggregatedBalanceSelector'; +const TRACE_TAGS = { controller: 'AssetsController' }; + +/** + * Record the aggregation timing as a single subspan under one parent span. + * + * These selectors are synchronous, so the work is already over by the time the + * spans are opened: both are backdated by `durationMs`, and the parent carries + * it as a numeric tag so MetaMask's Sentry adapter records the `duration_ms` + * measurement the dashboards chart. + * + * @param trace - Trace callback from the client. + * @param durationMs - Measured compute time in milliseconds. + * @param data - Additional span attributes, such as counts. + */ +function emitAggregatedBalanceTrace( + trace: TraceCallback, + durationMs: number, + data: Record, +): void { + const spanData = { duration_ms: durationMs, ...data }; + + trace( + { + name: TRACE_AGGREGATED_BALANCE, + data: spanData, + tags: { ...TRACE_TAGS, duration_ms: durationMs }, + startTime: performance.timeOrigin + performance.now() - durationMs, + }, + (parentContext) => { + emitTrace( + trace, + { parent: parentContext }, + { + name: TRACE_AGGREGATED_BALANCE_SELECTOR, + data: spanData, + tags: TRACE_TAGS, + }, + ); + return undefined; + }, + ).catch(() => { + // Telemetry failure must not break. + }); +} + export type EnabledNetworkMap = | Record> | undefined; @@ -523,20 +571,10 @@ function aggregateBalances( const info = getAssetInfo(assetInfoCache, assetId); uniqueNetworks.add(info.chainId); } - trace( - { - name: TRACE_AGGREGATED_BALANCE_SELECTOR, - data: { - duration_ms: durationMs, - asset_count: merged.size, - network_count: uniqueNetworks.size, - account_count: accountsToAggregate.length, - }, - tags: { controller: 'AssetsController' }, - }, - () => undefined, - ).catch(() => { - // Telemetry failure must not break. + emitAggregatedBalanceTrace(trace, durationMs, { + asset_count: merged.size, + network_count: uniqueNetworks.size, + account_count: accountsToAggregate.length, }); } @@ -721,9 +759,11 @@ export function calculateBalanceForAllWallets( enabledNetworkMap?: EnabledNetworkMap, trace?: TraceCallback, ): AllWalletsBalance { + const startTime = performance.now(); const userCurrency = getUserCurrency(assetsControllerState); const wallets: AllWalletsBalance['wallets'] = {}; let totalBalanceInUserCurrency = 0; + let groupCount = 0; type WalletWithGroups = { groups?: Record }; for (const [walletId, wallet] of Object.entries( @@ -738,12 +778,15 @@ export function calculateBalanceForAllWallets( const groups = (wallet as WalletWithGroups)?.groups ?? {}; for (const groupId of Object.keys(groups)) { + groupCount += 1; const accountIds = getAccountIdsForGroup(accountTreeState, groupId); + // `trace` is deliberately not forwarded: doing so opened one root span + // per account group and hit Sentry's rate limits. One span for the whole + // walk is emitted below instead. const { totalBalanceInFiat = 0 } = getAggregatedBalanceForAccountIds( assetsControllerState, accountIds, enabledNetworkMap, - trace, ); walletBalance.groups[groupId] = { @@ -759,6 +802,13 @@ export function calculateBalanceForAllWallets( totalBalanceInUserCurrency += walletBalance.totalBalanceInUserCurrency; } + if (trace) { + emitAggregatedBalanceTrace(trace, performance.now() - startTime, { + wallet_count: Object.keys(wallets).length, + group_count: groupCount, + }); + } + return { wallets, totalBalanceInUserCurrency, userCurrency }; } diff --git a/packages/assets-controller/src/traced.test.ts b/packages/assets-controller/src/traced.test.ts new file mode 100644 index 00000000000..0b0f887c444 --- /dev/null +++ b/packages/assets-controller/src/traced.test.ts @@ -0,0 +1,500 @@ +import type { TraceCallback, TraceRequest } from '@metamask/controller-utils'; + +import type { SpanHandle } from './traced.js'; +import { emitTrace, registerTracer, traced } from './traced.js'; + +/** + * A trace callback that records every request and hands each span a unique + * context object, so tests can assert nesting by identity. + */ +type Recorder = { + trace: TraceCallback; + requests: TraceRequest[]; + contexts: unknown[]; +}; + +/** + * Build a recording trace callback. + * + * @returns The callback plus the requests and contexts it has seen. + */ +function makeRecorder(): Recorder { + const requests: TraceRequest[] = []; + const contexts: unknown[] = []; + const trace = (async ( + request: TraceRequest, + fn?: (context?: unknown) => unknown, + ) => { + requests.push(request); + const context = { spanId: `span-${requests.length}` }; + contexts.push(context); + return fn?.(context); + }) as TraceCallback; + return { trace, requests, contexts }; +} + +/** + * Find the requests with a given span name. + * + * @param requests - Requests seen by the recorder. + * @param name - Span name to filter by. + * @returns The matching requests. + */ +function withName(requests: TraceRequest[], name: string): TraceRequest[] { + return requests.filter((request) => request.name === name); +} + +/** + * A controller with two nested traced methods, modelling a public entry point + * that decides the root tracing policy and a private choke point below it. + */ +class DemoController { + readonly #label = 'demo'; + + #firstInitReported = false; + + readonly executionCounts: Record = {}; + + constructor(trace?: TraceCallback) { + registerTracer(this, trace); + } + + async getAssets(accounts: string[]): Promise { + const span: SpanHandle = { enabled: !this.#firstInitReported }; + this.#firstInitReported = true; + return this.#fullFetch(span, accounts); + } + + @traced({ + name: 'AssetsFetchPipeline', + tags: { controller: 'DemoController' }, + data: (accounts: string[]) => ({ account_count: accounts.length }), + }) + async #fullFetch(span: SpanHandle, accounts: string[]): Promise { + this.#count('fullFetch'); + // `span` is the handle injected by the decorator, carrying this method's + // own span context; passing it to a traced child nests the child. + return this.#enrich(span, accounts); + } + + @traced({ name: 'AssetsUpdateEnrichment' }) + async #enrich(_span: SpanHandle, accounts: string[]): Promise { + this.#count('enrich'); + if (accounts.includes('boom')) { + throw new Error('enrichment failed'); + } + return accounts.map((account) => `${account}:${this.#label}`); + } + + #count(key: string): void { + this.executionCounts[key] = (this.executionCounts[key] ?? 0) + 1; + } +} + +const LANE_SPAN_NAMES = { + fast: 'AssetsFetchPipeline', + background: 'AssetsBackgroundFetch', + update: 'AssetsUpdateEnrichment', +} as const; + +type Lane = keyof typeof LANE_SPAN_NAMES; +type PipelineOptions = { lane: Lane; sources: string[] }; +type PipelineResult = { balances: Record }; + +/** + * A controller whose single decorated choke point serves several lanes, each + * keeping its own span name and summary records. + */ +class PipelineController { + tracingEnabled = true; + + readonly #trace?: TraceCallback; + + constructor(trace?: TraceCallback) { + this.#trace = trace; + registerTracer(this, trace); + } + + async run( + lane: Lane, + sources: string[], + initial?: Record, + ): Promise { + const span: SpanHandle = { enabled: this.tracingEnabled }; + return initial === undefined + ? this.#executeMiddlewares(span, { lane, sources }) + : this.#executeMiddlewares(span, { lane, sources }, initial); + } + + @traced({ + name: (options: PipelineOptions) => LANE_SPAN_NAMES[options.lane], + tags: { controller: 'PipelineController' }, + data: (options: PipelineOptions) => ({ + lane: options.lane, + source_count: options.sources.length, + }), + summary: ( + result: PipelineResult, + elapsedMs: number, + options: PipelineOptions, + ) => + options.lane === 'fast' + ? { + name: 'AssetsFullFetch', + data: { + duration_ms: elapsedMs, + asset_count: Object.keys(result.balances).length, + }, + tags: { controller: 'PipelineController' }, + } + : undefined, + }) + async #executeMiddlewares( + span: SpanHandle, + options: PipelineOptions, + initialResponse: Record = {}, + ): Promise { + const balances = { ...initialResponse }; + for (const source of options.sources) { + balances[source] = 1; + emitTrace(this.#trace, span, { + name: 'AssetsDataSourceTiming', + data: { source, duration_ms: 5 }, + tags: { source }, + }); + } + return { balances }; + } +} + +describe('traced', () => { + describe('span creation and nesting', () => { + it('returns the work result and preserves access to private fields', async () => { + const { trace } = makeRecorder(); + const controller = new DemoController(trace); + + expect(await controller.getAssets(['acc1', 'acc2'])).toStrictEqual([ + 'acc1:demo', + 'acc2:demo', + ]); + }); + + it('emits a root span for the outermost traced method', async () => { + const { trace, requests } = makeRecorder(); + const controller = new DemoController(trace); + + await controller.getAssets(['acc1']); + + expect(requests).toHaveLength(2); + expect(requests[0]).toMatchObject({ + name: 'AssetsFetchPipeline', + parentContext: undefined, + tags: { controller: 'DemoController' }, + }); + }); + + it('passes the call arguments to the data extractor', async () => { + const { trace, requests } = makeRecorder(); + const controller = new DemoController(trace); + + await controller.getAssets(['acc1', 'acc2']); + + expect(requests[0]?.data).toStrictEqual({ account_count: 2 }); + }); + + it('nests a child span under the context injected into its parent', async () => { + const { trace, requests, contexts } = makeRecorder(); + const controller = new DemoController(trace); + + await controller.getAssets(['acc1']); + + expect(requests[1]).toMatchObject({ + name: 'AssetsUpdateEnrichment', + parentContext: contexts[0], + }); + }); + }); + + describe('when tracing is disabled for the call tree', () => { + it('silences every span below the entry point but still runs the work', async () => { + const { trace, requests } = makeRecorder(); + const controller = new DemoController(trace); + + await controller.getAssets(['acc1']); + const spansAfterFirstCall = requests.length; + const result = await controller.getAssets(['acc2']); + + expect(requests).toHaveLength(spansAfterFirstCall); + expect(result).toStrictEqual(['acc2:demo']); + }); + }); + + describe('when no tracer is registered', () => { + it('passes through to the work untouched', async () => { + const controller = new DemoController(undefined); + + expect(await controller.getAssets(['acc1'])).toStrictEqual(['acc1:demo']); + }); + }); + + describe('when the tracer fails', () => { + it('returns the work result if the tracer rejects after the work ran', async () => { + const trace = (async ( + _request: TraceRequest, + fn?: (context?: unknown) => unknown, + ) => { + if (fn) { + await fn({ id: 'parent' }); + throw new Error('telemetry failed'); + } + return undefined; + }) as TraceCallback; + const controller = new DemoController(trace); + + expect(await controller.getAssets(['acc1'])).toStrictEqual(['acc1:demo']); + expect(controller.executionCounts.enrich).toBe(1); + }); + + it('runs the work exactly once if the tracer rejects before invoking it', async () => { + const trace = (async () => { + throw new Error('telemetry failed'); + }) as TraceCallback; + const controller = new DemoController(trace); + + expect(await controller.getAssets(['acc1'])).toStrictEqual(['acc1:demo']); + expect(controller.executionCounts).toStrictEqual({ + fullFetch: 1, + enrich: 1, + }); + }); + }); + + describe('when the work throws', () => { + it('propagates the error rather than swallowing it as telemetry', async () => { + const { trace } = makeRecorder(); + const controller = new DemoController(trace); + + await expect(controller.getAssets(['boom'])).rejects.toThrow( + 'enrichment failed', + ); + }); + }); + + describe('as a single choke point for several lanes', () => { + it('accepts an object parameter alongside a trailing defaulted parameter', async () => { + const { trace } = makeRecorder(); + const controller = new PipelineController(trace); + + const result = await controller.run('fast', ['api', 'staked'], { + seed: 1, + }); + + expect(result.balances).toStrictEqual({ seed: 1, api: 1, staked: 1 }); + }); + + it('derives the span name from the call arguments', async () => { + const { trace, requests } = makeRecorder(); + const controller = new PipelineController(trace); + + await controller.run('fast', ['api']); + + expect(requests[0]).toMatchObject({ + name: 'AssetsFetchPipeline', + data: { lane: 'fast', source_count: 1 }, + }); + }); + + it('gives a different lane its own root span name', async () => { + const { trace, requests } = makeRecorder(); + const controller = new PipelineController(trace); + + await controller.run('background', ['rpc']); + + expect(requests[0]).toMatchObject({ + name: 'AssetsBackgroundFetch', + parentContext: undefined, + }); + }); + + it('nests records emitted from the method body under its own span', async () => { + const { trace, requests, contexts } = makeRecorder(); + const controller = new PipelineController(trace); + + await controller.run('fast', ['api', 'staked']); + + const timings = withName(requests, 'AssetsDataSourceTiming'); + expect(timings).toHaveLength(2); + expect( + timings.every((request) => request.parentContext === contexts[0]), + ).toBe(true); + }); + }); + + describe('summary records', () => { + it('derives them from the result and nests them under the span', async () => { + const { trace, requests, contexts } = makeRecorder(); + const controller = new PipelineController(trace); + + await controller.run('fast', ['api', 'staked']); + + expect(withName(requests, 'AssetsFullFetch')[0]).toMatchObject({ + parentContext: contexts[0], + data: { duration_ms: expect.any(Number), asset_count: 2 }, + tags: { duration_ms: expect.any(Number), asset_count: 2 }, + startTime: expect.any(Number), + }); + }); + + it('emits them while the span is still open', async () => { + // A record attached to an already-finished Sentry transaction is + // orphaned and dropped, so the parent must still be open at this point. + const openSpans: string[] = []; + const openWhenSummaryEmitted: string[][] = []; + const trace = (async ( + request: TraceRequest, + fn?: (context?: unknown) => unknown, + ) => { + if (request.name === 'AssetsFullFetch') { + openWhenSummaryEmitted.push([...openSpans]); + } + if (!fn) { + return undefined; + } + openSpans.push(request.name); + try { + return await fn({ spanId: request.name }); + } finally { + openSpans.pop(); + } + }) as TraceCallback; + + await new PipelineController(trace).run('fast', ['api']); + + expect(openWhenSummaryEmitted).toStrictEqual([['AssetsFetchPipeline']]); + }); + + it('emits nothing when the extractor returns undefined', async () => { + const { trace, requests } = makeRecorder(); + const controller = new PipelineController(trace); + + await controller.run('background', ['rpc']); + + expect(withName(requests, 'AssetsFullFetch')).toHaveLength(0); + }); + + it('is silenced along with the wrap span when tracing is disabled', async () => { + const { trace, requests } = makeRecorder(); + const controller = new PipelineController(trace); + controller.tracingEnabled = false; + + const result = await controller.run('fast', ['api']); + + expect(requests).toHaveLength(0); + expect(result.balances).toStrictEqual({ api: 1 }); + }); + }); +}); + +describe('emitTrace', () => { + it('copies numeric data fields into tags so Sentry records measurements', () => { + const { trace, requests } = makeRecorder(); + + emitTrace(trace, undefined, { + name: 'AssetsStateSize', + data: { asset_count: 12, source: 'api' }, + tags: { controller: 'AssetsController' }, + }); + + expect(requests[0]?.tags).toStrictEqual({ + controller: 'AssetsController', + asset_count: 12, + }); + }); + + it('backdates startTime by duration_ms so the span renders its real length', () => { + const { trace, requests } = makeRecorder(); + const before = performance.timeOrigin + performance.now(); + + emitTrace(trace, undefined, { + name: 'AssetsFullFetch', + data: { duration_ms: 250 }, + }); + + const startTime = requests[0]?.startTime as number; + expect(startTime).toBeGreaterThanOrEqual(before - 250); + expect(startTime).toBeLessThanOrEqual( + performance.timeOrigin + performance.now() - 250, + ); + }); + + it('omits startTime when there is no duration to backdate from', () => { + const { trace, requests } = makeRecorder(); + + emitTrace(trace, undefined, { name: 'AssetsSubscriptionError' }); + + expect(requests[0]).toStrictEqual({ + name: 'AssetsSubscriptionError', + data: {}, + tags: {}, + parentContext: undefined, + }); + }); + + it('nests the record under the handle parent', () => { + const { trace, requests } = makeRecorder(); + const parent = { spanId: 'parent' }; + + emitTrace(trace, { parent }, { name: 'AssetsDataSourceTiming' }); + + expect(requests[0]?.parentContext).toBe(parent); + }); + + it('does nothing without a tracer', () => { + expect(() => + emitTrace(undefined, undefined, { name: 'AssetsFullFetch' }), + ).not.toThrow(); + }); + + it('does nothing when the handle disables tracing', () => { + const { trace, requests } = makeRecorder(); + + emitTrace(trace, { enabled: false }, { name: 'AssetsFullFetch' }); + + expect(requests).toHaveLength(0); + }); + + it('swallows a tracer that throws synchronously', () => { + const trace = (() => { + throw new Error('telemetry failed'); + }) as TraceCallback; + + expect(() => + emitTrace(trace, undefined, { name: 'AssetsFullFetch' }), + ).not.toThrow(); + }); + + it('swallows a tracer that rejects', async () => { + const unhandled = jest.fn(); + process.once('unhandledRejection', unhandled); + const trace = (async () => { + throw new Error('telemetry failed'); + }) as TraceCallback; + + emitTrace(trace, undefined, { name: 'AssetsFullFetch' }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(unhandled).not.toHaveBeenCalled(); + process.off('unhandledRejection', unhandled); + }); +}); + +describe('registerTracer', () => { + it('leaves the instance untraced when given no callback', async () => { + const { requests } = makeRecorder(); + const controller = new DemoController(undefined); + + await controller.getAssets(['acc1']); + + expect(requests).toHaveLength(0); + }); +}); diff --git a/packages/assets-controller/src/traced.ts b/packages/assets-controller/src/traced.ts new file mode 100644 index 00000000000..be2f5c60363 --- /dev/null +++ b/packages/assets-controller/src/traced.ts @@ -0,0 +1,293 @@ +import type { TraceCallback, TraceContext } from '@metamask/controller-utils'; + +/** + * Decorator-based Sentry tracing. + * + * The protocol: every traced method takes a {@link SpanHandle} as its FIRST + * parameter. Callers pass their own handle; {@link traced} consumes it (using + * `.parent` to nest and `.enabled` to decide whether to emit at all) and + * injects a NEW handle carrying the freshly created span's context into the + * method body. Inside a traced method, `span` is therefore *your own* span — + * pass it to traced children or to {@link emitTrace} to nest them underneath. + * + * `enabled: false` short-circuits tracing and propagates down the entire call + * tree, so a single decision at a public entry point silences everything below + * it without conditionals sprinkled through the implementation. + * + * The handle is the first parameter rather than the last because methods with + * trailing optional or defaulted parameters make a trailing handle unreliable + * to locate at runtime. + */ + +/** + * A tracing handle threaded through a call tree. + */ +export type SpanHandle = { + /** Context of the enclosing span; children nest under it. */ + parent?: TraceContext; + /** When `false`, no spans are emitted anywhere below this point. */ + enabled?: boolean; +}; + +/** + * A zero-duration span used to record a set of measurements, as opposed to + * wrapping a unit of work. + */ +export type SummaryRecord = { + /** Span name visible in Sentry. */ + name: string; + /** Key-value pairs attached as span data. */ + data?: Record; + /** Key-value pairs used for Sentry filtering. */ + tags?: Record; +}; + +/** + * Decorator wrappers cannot read `#private` fields — private names are + * lexically scoped to the class body — so instances publish their tracer here + * instead. Module-private, therefore as encapsulated as a `#field`. + */ +const tracers = new WeakMap(); + +/** + * Associate a trace callback with an instance so that {@link traced} methods on + * it emit spans. Call this from the constructor. Instances registered without a + * callback stay untraced, and their decorated methods become pure passthroughs. + * + * @param instance - The instance whose decorated methods should be traced. + * @param trace - Trace callback supplied by the client, if any. + */ +export function registerTracer( + instance: object, + trace: TraceCallback | undefined, +): void { + if (trace) { + tracers.set(instance, trace); + } +} + +/** + * Emit a fire-and-forget record span. The handle internalizes both whether to + * trace at all (`span.enabled`) and where to nest (`span.parent`). + * + * Numeric `data` fields are copied into `tags` because MetaMask's Sentry + * adapter promotes numeric tags to measurements, which is what dashboard + * widgets chart (`tags[duration_ms,number]`). When `data.duration_ms` is + * present, `startTime` is backdated by that amount so the span's rendered + * duration matches the work that was measured. + * + * @param trace - Trace callback supplied by the client, if any. + * @param span - Handle of the enclosing span, if any. + * @param record - The span to record. + */ +export function emitTrace( + trace: TraceCallback | undefined, + span: SpanHandle | undefined, + record: SummaryRecord, +): void { + if (!trace || span?.enabled === false) { + return; + } + + const data = record.data ?? {}; + const numericFromData: Record = {}; + for (const [key, value] of Object.entries(data)) { + if (typeof value === 'number') { + numericFromData[key] = value; + } + } + + const durationMs = + typeof data.duration_ms === 'number' ? data.duration_ms : undefined; + const startTime = + durationMs === undefined + ? undefined + : performance.timeOrigin + performance.now() - durationMs; + + try { + trace( + { + name: record.name, + data, + tags: { ...record.tags, ...numericFromData }, + parentContext: span?.parent, + ...(startTime === undefined ? {} : { startTime }), + }, + () => undefined, + )?.catch(() => { + // Telemetry failure must not break. + }); + } catch { + // Telemetry failure must not break. + } +} + +/** + * Options for {@link traced}. + * + * The extractors are typed against a PREFIX of the decorated method's + * arguments (excluding the span handle), so an extractor only has to type the + * parameters it actually reads even when the method has further trailing or + * defaulted parameters. + */ +export type TracedOptions< + ExtractorArgs extends readonly unknown[], + ExtractorResult, +> = { + /** Span name, or a function deriving it from the call's arguments. */ + name: string | ((...args: ExtractorArgs) => string); + /** Key-value pairs used for Sentry filtering. */ + tags?: Record; + /** Derives span data from the call's arguments. */ + data?: (...args: ExtractorArgs) => Record; + /** + * Derives record spans from the result. They are emitted while the span is + * still open so that they nest under it: a child attached to an + * already-finished Sentry transaction is orphaned and dropped. + */ + summary?: ( + result: ExtractorResult, + elapsedMs: number, + ...args: ExtractorArgs + ) => SummaryRecord | SummaryRecord[] | undefined; +}; + +/** + * Wrap an async method in a Sentry span. + * + * The decorated method must take a {@link SpanHandle} as its first parameter; + * see the module documentation for the protocol. + * + * Telemetry never breaks the work: if the tracer rejects after the work ran, + * the work's result is returned; if it rejects before invoking the work, the + * work runs once without a span. Errors thrown by the work itself propagate. + * + * @param options - See {@link TracedOptions}. + * @returns A method decorator. + */ +export function traced< + ExtractorArgs extends readonly unknown[] = readonly unknown[], + ExtractorResult = unknown, +>(options: TracedOptions) { + return function decorate< + This extends object, + Args extends readonly [...ExtractorArgs, ...unknown[]], + Return extends ExtractorResult, + >( + // The handle is declared required here so that a method may declare it + // either way: `span: SpanHandle | undefined` still satisfies this + // contravariantly, while `span: SpanHandle` forces every call site to + // state whether it wants tracing. + target: (this: This, span: SpanHandle, ...args: Args) => Promise, + _context: ClassMethodDecoratorContext< + This, + (this: This, span: SpanHandle, ...args: Args) => Promise + >, + ) { + // Extractors are typed against a prefix of the method's arguments; JS + // ignores extra arguments, so calling them with all of them is safe. + const getName = (args: Args): string => + typeof options.name === 'function' + ? (options.name as (...a: readonly unknown[]) => string)(...args) + : options.name; + + const getData = ( + args: Args, + ): Record | undefined => + options.data + ? ( + options.data as ( + ...a: readonly unknown[] + ) => Record + )(...args) + : undefined; + + const getSummaryRecords = ( + value: Return, + elapsedMs: number, + args: Args, + ): SummaryRecord[] => { + if (!options.summary) { + return []; + } + const produced = ( + options.summary as ( + result: unknown, + elapsed: number, + ...a: readonly unknown[] + ) => SummaryRecord | SummaryRecord[] | undefined + )(value, elapsedMs, ...args); + if (produced === undefined) { + return []; + } + return Array.isArray(produced) ? produced : [produced]; + }; + + return async function tracedMethod( + this: This, + span: SpanHandle, + ...args: Args + ): Promise { + const trace = tracers.get(this); + if (!trace || span?.enabled === false) { + return target.call(this, span, ...args); + } + + let outcome: + | { ok: true; value: Return } + | { ok: false; error: unknown } + | undefined; + + try { + await trace( + { + name: getName(args), + parentContext: span?.parent, + data: getData(args), + tags: options.tags, + }, + async (context) => { + const startedMs = performance.now(); + try { + const value = await target.call( + this, + { parent: context, enabled: span?.enabled }, + ...args, + ); + outcome = { ok: true, value }; + + try { + const elapsedMs = performance.now() - startedMs; + for (const record of getSummaryRecords( + value, + elapsedMs, + args, + )) { + emitTrace(trace, { parent: context }, record); + } + } catch { + // Summary telemetry must not break the work. + } + + return value; + } catch (error) { + outcome = { ok: false, error }; + throw error; + } + }, + ); + } catch { + // Telemetry failure must not break the work — unless the work failed. + } + + if (outcome === undefined) { + // The tracer failed before invoking the work; run it without a span. + return target.call(this, { enabled: span?.enabled }, ...args); + } + if (outcome.ok) { + return outcome.value; + } + throw outcome.error; + }; + }; +}