From 6101c32bd992c871497203bcb6d509ba4610fea5 Mon Sep 17 00:00:00 2001
From: dan <dan.abramov@gmail.com>
Date: Thu, 18 Apr 2024 16:26:05 +0100
Subject: [PATCH] [Statsig] Prefetch configs for other accounts (#3607)

* Poll both current and other accounts

* Make createStatsigOptions a function

* Pass prefetchUsers with the initial request

* Add initializeCalled check

* Be resilient to object identity changes

* Decrease poll interval to 1 minute
---
 src/lib/statsig/statsig.tsx | 67 +++++++++++++++++++++++++------------
 1 file changed, 46 insertions(+), 21 deletions(-)

diff --git a/src/lib/statsig/statsig.tsx b/src/lib/statsig/statsig.tsx
index 36f030e3c6..151b365d34 100644
--- a/src/lib/statsig/statsig.tsx
+++ b/src/lib/statsig/statsig.tsx
@@ -8,6 +8,7 @@ import {logger} from '#/logger'
 import {isWeb} from '#/platform/detection'
 import {IS_TESTFLIGHT} from 'lib/app-info'
 import {useSession} from '../../state/session'
+import {useNonReactiveCallback} from '../hooks/useNonReactiveCallback'
 import {LogEvents} from './events'
 import {Gate} from './gates'
 
@@ -34,19 +35,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<
@@ -160,9 +165,25 @@ AppState.addEventListener('change', (state: AppStateStatus) => {
 })
 
 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 +191,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}>