Merge remote-tracking branch 'origin/main' into halo/base

* origin/main:
  [Statsig] Slightly block the UI on gates (#3608)
  [Statsig] Prefetch configs for other accounts (#3607)

src/lib/statsig/statsig.tsx

@@ -8,6 +8,8 @@ import {logger} from '#/logger'
 import {isWeb} from '#/platform/detection'
 import {IS_TESTFLIGHT} from 'lib/app-info'
 import {useSession} from '../../state/session'
+import {timeout} from '../async/timeout'
+import {useNonReactiveCallback} from '../hooks/useNonReactiveCallback'
 import {LogEvents} from './events'
 import {Gate} from './gates'
 
@@ -34,19 +36,23 @@ if (isWeb && typeof window !== 'undefined') {
 
 export type {LogEvents}
 
-const statsigOptions = {
-  environment: {
-    tier:
-      process.env.NODE_ENV === 'development'
-        ? 'development'
-        : IS_TESTFLIGHT
-        ? 'staging'
-        : 'production',
-  },
-  // Don't block on waiting for network. The fetched config will kick in on next load.
-  // This ensures the UI is always consistent and doesn't update mid-session.
-  // Note this makes cold load (no local storage) and private mode return `false` for all gates.
-  initTimeoutMs: 1,
+function createStatsigOptions(prefetchUsers: StatsigUser[]) {
+  return {
+    environment: {
+      tier:
+        process.env.NODE_ENV === 'development'
+          ? 'development'
+          : IS_TESTFLIGHT
+          ? 'staging'
+          : 'production',
+    },
+    // Don't block on waiting for network. The fetched config will kick in on next load.
+    // This ensures the UI is always consistent and doesn't update mid-session.
+    // Note this makes cold load (no local storage) and private mode return `false` for all gates.
+    initTimeoutMs: 1,
+    // Get fresh flags for other accounts as well, if any.
+    prefetchUsers,
+  }
 }
 
 type FlatJSONRecord = Record<
@@ -159,10 +165,51 @@ AppState.addEventListener('change', (state: AppStateStatus) => {
   }
 })
 
+export async function tryFetchGates(
+  did: string,
+  strategy: 'prefer-low-latency' | 'prefer-fresh-gates',
+) {
+  try {
+    let timeoutMs = 250 // Don't block the UI if we can't do this fast.
+    if (strategy === 'prefer-fresh-gates') {
+      // Use this for less common operations where the user would be OK with a delay.
+      timeoutMs = 1500
+    }
+    // Note: This condition is currently false the very first render because
+    // Statsig has not initialized yet. In the future, we can fix this by
+    // doing the initialization ourselves instead of relying on the provider.
+    if (Statsig.initializeCalled()) {
+      await Promise.race([
+        timeout(timeoutMs),
+        Statsig.prefetchUsers([toStatsigUser(did)]),
+      ])
+    }
+  } catch (e) {
+    // Don't leak errors to the calling code, this is meant to be always safe.
+    console.error(e)
+  }
+}
+
 export function Provider({children}: {children: React.ReactNode}) {
-  const {currentAccount} = useSession()
+  const {currentAccount, accounts} = useSession()
   const did = currentAccount?.did
   const currentStatsigUser = React.useMemo(() => toStatsigUser(did), [did])
+
+  const otherDidsConcatenated = accounts
+    .map(account => account.did)
+    .filter(accountDid => accountDid !== did)
+    .join(' ') // We're only interested in DID changes.
+  const otherStatsigUsers = React.useMemo(
+    () => otherDidsConcatenated.split(' ').map(toStatsigUser),
+    [otherDidsConcatenated],
+  )
+  const statsigOptions = React.useMemo(
+    () => createStatsigOptions(otherStatsigUsers),
+    [otherStatsigUsers],
+  )
+
+  // Have our own cache in front of Statsig.
+  // This ensures the results remain stable until the active DID changes.
   const [gateCache, setGateCache] = React.useState(() => new Map())
   const [prevDid, setPrevDid] = React.useState(did)
   if (did !== prevDid) {
@@ -170,15 +217,19 @@ export function Provider({children}: {children: React.ReactNode}) {
     setGateCache(new Map())
   }
 
-  React.useEffect(() => {
-    function refresh() {
-      // This will not affect the current session.
-      // Statsig will put the results into local storage and we'll pick it up on next load.
-      Statsig.updateUser(currentStatsigUser)
+  // Periodically poll Statsig to get the current rule evaluations for all stored accounts.
+  // These changes are prefetched and stored, but don't get applied until the active DID changes.
+  // This ensures that when you switch an account, it already has fresh results by then.
+  const handleIntervalTick = useNonReactiveCallback(() => {
+    if (Statsig.initializeCalled()) {
+      // Note: Only first five will be taken into account by Statsig.
+      Statsig.prefetchUsers([currentStatsigUser, ...otherStatsigUsers])
     }
-    const id = setInterval(refresh, 3 * 60e3 /* 3 min */)
+  })
+  React.useEffect(() => {
+    const id = setInterval(handleIntervalTick, 60e3 /* 1 min */)
     return () => clearInterval(id)
-  }, [currentStatsigUser])
+  }, [handleIntervalTick])
 
   return (
     <GateCache.Provider value={gateCache}>

src/state/session/index.tsx

@@ -9,7 +9,7 @@ import {jwtDecode} from 'jwt-decode'
 import {track} from '#/lib/analytics/analytics'
 import {networkRetry} from '#/lib/async/retry'
 import {IS_TEST_USER} from '#/lib/constants'
-import {logEvent, LogEvents} from '#/lib/statsig/statsig'
+import {logEvent, LogEvents, tryFetchGates} from '#/lib/statsig/statsig'
 import {hasProp} from '#/lib/type-guards'
 import {logger} from '#/logger'
 import {isWeb} from '#/platform/detection'
@@ -243,6 +243,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       if (!agent.session) {
         throw new Error(`session: createAccount failed to establish a session`)
       }
+      const fetchingGates = tryFetchGates(
+        agent.session.did,
+        'prefer-fresh-gates',
+      )
 
       const deactivated = isSessionDeactivated(agent.session.accessJwt)
       if (!deactivated) {
@@ -283,6 +287,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       )
 
       __globalAgent = agent
+      await fetchingGates
       upsertAccount(account)
 
       logger.debug(`session: created account`, {}, logger.DebugContext.session)
@@ -303,6 +308,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       if (!agent.session) {
         throw new Error(`session: login failed to establish a session`)
       }
+      const fetchingGates = tryFetchGates(
+        agent.session.did,
+        'prefer-fresh-gates',
+      )
 
       const account: SessionAccount = {
         service: agent.service.toString(),
@@ -330,6 +339,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       __globalAgent = agent
       // @ts-ignore
       if (IS_DEV && isWeb) window.agent = agent
+      await fetchingGates
       upsertAccount(account)
 
       logger.debug(`session: logged in`, {}, logger.DebugContext.session)
@@ -362,6 +372,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   const initSession = React.useCallback<ApiContext['initSession']>(
     async account => {
       logger.debug(`session: initSession`, {}, logger.DebugContext.session)
+      const fetchingGates = tryFetchGates(account.did, 'prefer-low-latency')
 
       const agent = new BskyAgent({
         service: account.service,
@@ -406,6 +417,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
 
         agent.session = prevSession
         __globalAgent = agent
+        await fetchingGates
         upsertAccount(account)
 
         if (prevSession.deactivated) {
@@ -442,6 +454,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         try {
           const freshAccount = await resumeSessionWithFreshAccount()
           __globalAgent = agent
+          await fetchingGates
           upsertAccount(freshAccount)
         } catch (e) {
           /*