feat: perf event handlers (#553)

* build: add jest-extended

* feat: type new perf handlers

* feat: onSwapQuote perf handler

* feat: onSwapSend/onWrapSend perf handlers

* feat: onTokenAllowance/onPermit2Allowance perf handlers

* fix: nits

* fix: type unimplemented error

* fix: borked test

* fix: borked typing

* fix: more borked typing

jest.config.js

@@ -11,5 +11,5 @@ module.exports = {
     '@uniswap/conedison/provider': '@uniswap/conedison/dist/provider/index.js',
   },
   setupFiles: ['<rootDir>/test/setup.ts'],
-  setupFilesAfterEnv: ['<rootDir>/test/setup-jest.ts'],
+  setupFilesAfterEnv: ['jest-extended/all', '<rootDir>/test/setup-jest.ts'],
 }

package.json

@@ -191,6 +191,7 @@
     "inter-ui": "^3.13.1",
     "jest": "^27.5.1",
     "jest-environment-hardhat": "^1.1.4",
+    "jest-extended": "^3.2.4",
     "jest-fetch-mock": "^3.0.3",
     "jest-styled-components": "^7.0.5",
     "json-loader": "^0.5.7",

src/components/Swap/SwapActionButton/SwapButton.tsx

@@ -14,8 +14,6 @@ import { useUniversalRouterSwapCallback } from 'hooks/useUniversalRouter'
 import { useAtomValue } from 'jotai/utils'
 import { useCallback, useEffect, useState } from 'react'
 import { feeOptionsAtom, Field, swapEventHandlersAtom } from 'state/swap'
-import { TransactionType } from 'state/transactions'
-import invariant from 'tiny-invariant'
 
 import ActionButton from '../../ActionButton'
 import { SummaryDialog } from '../Summary'
@@ -76,18 +74,9 @@ export default function SwapButton({ disabled }: { disabled: boolean }) {
 
         // Set the block containing the response to the oldest valid block to ensure that the
         // completed trade's impact is reflected in future fetched trades.
-        response.wait(1).then((receipt) => {
+        response.response.wait(1).then((receipt) => {
           setOldestValidBlock(receipt.blockNumber)
         })
-
-        invariant(trade)
-        return {
-          type: TransactionType.SWAP,
-          response,
-          tradeType: trade.tradeType,
-          trade,
-          slippageTolerance: slippage.allowed,
-        }
       })
 
       // Only close the review modal if the swap submitted (ie no-throw).
@@ -97,7 +86,7 @@ export default function SwapButton({ disabled }: { disabled: boolean }) {
     } catch (e) {
       throwAsync(e)
     }
-  }, [onSubmit, setOldestValidBlock, slippage.allowed, swapCallback, throwAsync, trade])
+  }, [onSubmit, setOldestValidBlock, swapCallback, throwAsync])
 
   const onReviewSwapClick = useConditionalHandler(useAtomValue(swapEventHandlersAtom).onReviewSwapClick)
   const collapseToolbar = useCollapseToolbar()

src/components/Swap/SwapActionButton/WrapButton.tsx

@@ -1,12 +1,10 @@
 import { Trans } from '@lingui/macro'
-import { CurrencyAmount } from '@uniswap/sdk-core'
 import useWrapCallback from 'hooks/swap/useWrapCallback'
 import useNativeCurrency from 'hooks/useNativeCurrency'
 import useTokenColorExtraction from 'hooks/useTokenColorExtraction'
 import { Spinner } from 'icons'
 import { useCallback, useEffect, useMemo, useState } from 'react'
 import { TransactionType } from 'state/transactions'
-import invariant from 'tiny-invariant'
 
 import ActionButton from '../../ActionButton'
 import useOnSubmit from './useOnSubmit'
@@ -30,20 +28,13 @@ export default function WrapButton({ disabled }: { disabled: boolean }) {
   const onWrap = useCallback(async () => {
     setIsPending(true)
     try {
-      await onSubmit(async () => {
-        const response = await wrapCallback()
-        if (!response) return
-
-        invariant(wrapType !== undefined) // if response is valid, then so is wrapType
-        const amount = CurrencyAmount.fromRawAmount(native, response.value?.toString() ?? '0')
-        return { response, type: wrapType, amount }
-      })
+      await onSubmit(wrapCallback)
     } catch (e) {
       console.error(e) // ignore error
     } finally {
       setIsPending(false)
     }
-  }, [native, onSubmit, wrapCallback, wrapType])
+  }, [onSubmit, wrapCallback])
 
   const actionProps = useMemo(
     () =>

src/cosmos/EventFeed.tsx

@@ -31,20 +31,26 @@ export const HANDLERS: (keyof SwapEventHandlers | keyof TransactionEventHandlers
   'onError',
   'onExpandSwapDetails',
   'onInitialSwapQuote',
-  'onSwapApprove',
+  'onPermit2Allowance',
   'onReviewSwapClick',
+  'onRouterPreferenceChange',
   'onSettingsReset',
   'onSlippageChange',
   'onSubmitSwapClick',
+  'onSwapApprove',
   'onSwapPriceUpdateAck',
+  'onSwapQuote',
+  'onSwapSend',
   'onSwitchChain',
   'onSwitchTokens',
+  'onTokenAllowance',
   'onTokenChange',
   'onTokenSelectorClick',
   'onTransactionDeadlineChange',
   'onTxFail',
   'onTxSubmit',
   'onTxSuccess',
+  'onWrapSend',
 ]
 
 export interface Event {

src/hooks/swap/useSwapCallback.tsx

@@ -1,5 +1,4 @@
 import { BigNumber } from '@ethersproject/bignumber'
-import { TransactionResponse } from '@ethersproject/providers'
 import { Trans } from '@lingui/macro'
 import { Percent } from '@uniswap/sdk-core'
 import { FeeOptions } from '@uniswap/v3-sdk'
@@ -9,6 +8,7 @@ import { SignatureData } from 'hooks/usePermit'
 import { useSwapCallArguments } from 'hooks/useSwapCallArguments'
 import { ReactNode, useMemo } from 'react'
 import { InterfaceTrade } from 'state/routing/types'
+import { SwapTransactionInfo, TransactionType } from 'state/transactions'
 
 import useSendSwapTransaction from './useSendSwapTransaction'
 
@@ -20,7 +20,7 @@ export enum SwapCallbackState {
 
 interface UseSwapCallbackReturns {
   state: SwapCallbackState
-  callback?: () => Promise<TransactionResponse>
+  callback?: () => Promise<SwapTransactionInfo>
   error?: ReactNode
 }
 interface UseSwapCallbackArgs {
@@ -71,7 +71,13 @@ export function useSwapCallback({
 
     return {
       state: SwapCallbackState.VALID,
-      callback: async () => callback(),
+      callback: async () => ({
+        type: TransactionType.SWAP,
+        response: await callback(),
+        tradeType: trade.tradeType,
+        trade,
+        slippageTolerance: allowedSlippage,
+      }),
     }
-  }, [trade, provider, account, chainId, callback, recipient, recipientAddressOrName])
+  }, [trade, provider, account, chainId, callback, recipient, recipientAddressOrName, allowedSlippage])
 }

src/hooks/swap/useWrapCallback.test.ts

@@ -0,0 +1,49 @@
+import { parseEther } from '@ethersproject/units'
+import { CurrencyAmount } from '@uniswap/sdk-core'
+import { SupportedChainId } from 'constants/chains'
+import { ExtendedEther } from 'constants/tokens'
+import { Field, stateAtom, swapEventHandlersAtom } from 'state/swap'
+import { TransactionType } from 'state/transactions'
+import { renderHook, waitFor } from 'test'
+
+import useWrapCallback from './useWrapCallback'
+
+const ETH = ExtendedEther.onChain(SupportedChainId.MAINNET)
+const WETH = ETH.wrapped
+const AMOUNT = CurrencyAmount.fromRawAmount(ETH, parseEther('1').toString())
+const WRAP_TRANSACTION_INFO = {
+  response: expect.any(Object),
+  type: TransactionType.WRAP,
+  amount: AMOUNT,
+}
+
+describe('useWrapCallback', () => {
+  it('sends wrap to wallet', async () => {
+    const { result } = renderHook(() => useWrapCallback(), {
+      initialAtomValues: [[stateAtom, { amount: '1', [Field.INPUT]: ETH, [Field.OUTPUT]: WETH }]],
+    })
+    expect(result.current.callback).toBeInstanceOf(Function)
+
+    const info = await waitFor(async () => {
+      const info = await result.current.callback()
+      expect(info).toBeDefined()
+      return info
+    })
+    expect(info).toEqual(WRAP_TRANSACTION_INFO)
+  })
+
+  it('triggers onWrapSend', async () => {
+    const onWrapSend = jest.fn()
+    const { result } = renderHook(() => useWrapCallback(), {
+      initialAtomValues: [
+        [stateAtom, { amount: '1', [Field.INPUT]: ETH, [Field.OUTPUT]: WETH }],
+        [swapEventHandlersAtom, { onWrapSend }],
+      ],
+    })
+    await waitFor(async () => {
+      await expect(result.current.callback()).resolves.toBeDefined()
+    })
+    expect(onWrapSend).toHaveBeenLastCalledWith(expect.objectContaining({ amount: AMOUNT }), expect.any(Promise))
+    await expect(onWrapSend.mock.calls.slice(-1)[0][1]).resolves.toEqual(WRAP_TRANSACTION_INFO)
+  })
+})

src/hooks/swap/useWrapCallback.tsx

@@ -1,15 +1,15 @@
-import { ContractTransaction } from '@ethersproject/contracts'
 import { useWeb3React } from '@web3-react/core'
 import { WRAPPED_NATIVE_CURRENCY } from 'constants/tokens'
 import { useWETHContract } from 'hooks/useContract'
+import { usePerfEventHandler } from 'hooks/usePerfEventHandler'
 import { useAtomValue } from 'jotai/utils'
-import { useMemo } from 'react'
+import { useCallback, useMemo } from 'react'
 import { Field, swapAtom } from 'state/swap'
-import { TransactionType } from 'state/transactions'
+import { TransactionType, UnwrapTransactionInfo, WrapTransactionInfo } from 'state/transactions'
 import tryParseCurrencyAmount from 'utils/tryParseCurrencyAmount'
 
 interface UseWrapCallbackReturns {
-  callback: () => Promise<ContractTransaction | void>
+  callback: () => Promise<WrapTransactionInfo | UnwrapTransactionInfo | void>
   type?: TransactionType.WRAP | TransactionType.UNWRAP
 }
 
@@ -43,19 +43,28 @@ export default function useWrapCallback(): UseWrapCallbackReturns {
     [inputCurrency, amount]
   )
 
-  const callback = useMemo(() => {
-    return async () => {
-      if (!parsedAmountIn) return
-      switch (wrapType) {
-        case TransactionType.WRAP:
-          return wrappedNativeCurrencyContract?.deposit({ value: `0x${parsedAmountIn.quotient.toString(16)}` })
-        case TransactionType.UNWRAP:
-          return wrappedNativeCurrencyContract?.withdraw(`0x${parsedAmountIn.quotient.toString(16)}`)
-        case undefined:
-          return undefined
-      }
+  const wrapCallback = useCallback(async (): Promise<WrapTransactionInfo | UnwrapTransactionInfo | void> => {
+    if (!parsedAmountIn || !wrappedNativeCurrencyContract) return
+    switch (wrapType) {
+      case TransactionType.WRAP:
+        return {
+          response: await wrappedNativeCurrencyContract.deposit({ value: `0x${parsedAmountIn.quotient.toString(16)}` }),
+          type: TransactionType.WRAP,
+          amount: parsedAmountIn,
+        }
+      case TransactionType.UNWRAP:
+        return {
+          response: await wrappedNativeCurrencyContract.withdraw(`0x${parsedAmountIn.quotient.toString(16)}`),
+          type: TransactionType.WRAP,
+          amount: parsedAmountIn,
+        }
+      case undefined:
+        return undefined
     }
   }, [parsedAmountIn, wrappedNativeCurrencyContract, wrapType])
 
+  const args = useMemo(() => parsedAmountIn && { amount: parsedAmountIn }, [parsedAmountIn])
+  const callback = usePerfEventHandler('onWrapSend', args, wrapCallback)
+
   return useMemo(() => ({ callback, type: wrapType }), [callback, wrapType])
 }

src/hooks/usePerfEventHandler.test.ts

@@ -0,0 +1,55 @@
+import { TradeResult } from 'state/routing/types'
+import { swapEventHandlersAtom } from 'state/swap'
+import { PerfEventHandlers } from 'state/swap/perf'
+import { renderHook } from 'test'
+
+import { usePerfEventHandler } from './usePerfEventHandler'
+
+describe('usePerfEventHandler', () => {
+  const onSwapQuote = jest.fn()
+  const callback = jest.fn()
+  const tradeResult = {} as TradeResult
+
+  beforeEach(() => {
+    onSwapQuote.mockReset()
+    callback.mockReset().mockReturnValue(tradeResult)
+  })
+
+  describe('with a perfHandler', () => {
+    describe('with args', () => {
+      it('returns a callback returning the event, wrapped by the perfHandler', async () => {
+        const args = {} as Parameters<NonNullable<PerfEventHandlers['onSwapQuote']>>[0]
+        const { result } = renderHook(() => usePerfEventHandler('onSwapQuote', args, callback), {
+          initialAtomValues: [[swapEventHandlersAtom, { onSwapQuote }]],
+        })
+        expect(result.current).toBeInstanceOf(Function)
+        await expect(result.current()).resolves.toBe(tradeResult)
+
+        // The execution of the callback should be deferred until after the perfHandler has executed.
+        // This ensures that the perfHandler can capture the beginning of the callback's execution.
+        expect(onSwapQuote).toHaveBeenCalledBefore(callback)
+        expect(onSwapQuote).toHaveBeenCalledWith(args, expect.any(Promise))
+        expect(onSwapQuote.mock.calls[0][1]).resolves.toBe(tradeResult)
+      })
+    })
+
+    describe('without args', () => {
+      it('returns a callback returning the event, without calling perfHandler', async () => {
+        const { result } = renderHook(() => usePerfEventHandler('onSwapQuote', undefined, callback), {
+          initialAtomValues: [[swapEventHandlersAtom, { onSwapQuote }]],
+        })
+        expect(result.current).toBeInstanceOf(Function)
+        await expect(result.current()).resolves.toBe(tradeResult)
+        expect(onSwapQuote).not.toHaveBeenCalled()
+      })
+    })
+  })
+
+  describe('without a perfHandler', () => {
+    it('returns a callback returning the event', async () => {
+      const { result } = renderHook(() => usePerfEventHandler('onSwapQuote', undefined, callback))
+      expect(result.current).toBeInstanceOf(Function)
+      await expect(result.current()).resolves.toBe(tradeResult)
+    })
+  })
+})

src/hooks/usePerfEventHandler.ts

@@ -0,0 +1,27 @@
+import { useAtomValue } from 'jotai/utils'
+import { useCallback } from 'react'
+import { swapEventHandlersAtom } from 'state/swap'
+import { PerfEventHandlers } from 'state/swap/perf'
+
+/**
+ * PerfEventHandlers all take two arguments: args and event.
+ * This wraps those arguments so that the handler is called before the event is executed, for more accurate instrumentation.
+ */
+export function usePerfEventHandler<
+  Key extends keyof PerfEventHandlers,
+  Params extends Parameters<NonNullable<PerfEventHandlers[Key]>>,
+  Args extends Params[0],
+  Event extends Awaited<Params[1]>,
+  Handler extends PerfEventHandlers[Key] & ((args: Args, event: Promise<Event>) => void)
+>(name: Key, args: Args | undefined, callback: () => Promise<Event>): () => Promise<Event> {
+  const perfHandler = useAtomValue(swapEventHandlersAtom)[name] as Handler
+  return useCallback(() => {
+    // Use Promise.resolve().then to defer the execution of the callback until after the perfHandler has executed.
+    // This ensures that the perfHandler can capture the beginning of the callback's execution.
+    const event = Promise.resolve().then(callback)
+    if (args) {
+      perfHandler?.(args, event)
+    }
+    return event
+  }, [args, callback, perfHandler])
+}

src/hooks/usePermitAllowance.test.ts

@@ -9,6 +9,7 @@ import { UNI } from 'constants/tokens'
 import { useSingleCallResult } from 'hooks/multicall'
 import { useContract } from 'hooks/useContract'
 import ms from 'ms.macro'
+import { swapEventHandlersAtom } from 'state/swap'
 import { renderHook, waitFor } from 'test'
 
 import { usePermitAllowance, useUpdatePermitAllowance } from './usePermitAllowance'
@@ -120,6 +121,20 @@ describe('useUpdatePermitAllowance', () => {
     })
   })
 
+  it('triggers onPermit2Allowance', async () => {
+    const onPermit2Allowance = jest.fn()
+    const onPermitSignature = jest.fn()
+    const { result } = renderHook(() => useUpdatePermitAllowance(TOKEN, SPENDER, NONCE, onPermitSignature), {
+      initialAtomValues: [[swapEventHandlersAtom, { onPermit2Allowance }]],
+    })
+    await waitFor(() => result.current())
+    expect(onPermit2Allowance).toHaveBeenLastCalledWith(
+      expect.objectContaining({ token: TOKEN, spender: SPENDER }),
+      expect.any(Promise)
+    )
+    await expect(onPermit2Allowance.mock.calls.slice(-1)[0][1]).resolves.toBe(undefined)
+  })
+
   it('rejects on failure', async () => {
     const onPermitSignature = jest.fn()
     const { result } = renderHook(() => useUpdatePermitAllowance(TOKEN, SPENDER, NONCE, onPermitSignature))